model.py

import instructor
from openai import OpenAI

# import torch
# from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# # Set the cache directory
# cache_dir = "model"  # Replace with your desired folder path

# # Load the tokenizer and model with the specified cache directory
# tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct", cache_dir=cache_dir)
# model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B-Instruct", cache_dir=cache_dir)

# # Create the pipeline
# pipe = pipeline(
#     "text-generation",
#     model=model,
#     tokenizer=tokenizer,
#     torch_dtype=torch.bfloat16,
#     device_map="auto"
# )


llm = instructor.from_openai(
    OpenAI(
        base_url="http://localhost:11434/v1",
        api_key="ollama",
    ),
    mode=instructor.Mode.JSON,
)