Team Ai
Apppublic

salverz/llm-document-parser

sourceHugging Faceapache-2.0updated 1y agoView on Hugging Face
0likes
instructor_llm.py53 linesDownload Raw Back to llm_document_parser
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