salverz/llm-document-parser
0
1import instructor2from openai import OpenAI3from pydantic import BaseModel4from typing import Type5 6import ollama7 8def pull_ollama_model(model: str):9 """10 Pull a model from ollama if it is not already downloaded11 """12 if not model.__contains__(":"):13 model += ":latest"14 15 for downloaded_model in ollama.list()["models"]:16 if downloaded_model['model']== model:17 print(f"Model {downloaded_model['model']} is installed")18 return19 20 print(f"Model {model} is not installed")21 print(f"Downloading {model} model...")22 ollama.pull(model)23 24def extract_json_data_using_ollama_llm(prompt: str, text_data: str, ollama_model: str, response_model: Type[BaseModel]) -> str:25 """26 Pass prompt and data into an ollama LLM using instructor27 """28 client = instructor.from_openai(29 OpenAI(30 base_url="http://localhost:11434/v1",31 api_key="ollama"32 ),33 mode=instructor.Mode.JSON34 )35 36 resp = client.chat.completions.create(37 model=ollama_model,38 messages=[39 {40 'role': 'system',41 'content': prompt42 },43 {44 'role': 'user',45 'content': text_data46 },47 ],48 response_model=response_model,49 max_retries=350 )51 52 return resp.model_dump_json(indent=4)53 