src/model_selection.py

from langchain_openai import ChatOpenAI
from langchain_huggingface import ChatHuggingFace, HuggingFacePipeline
import os 
from dotenv import load_dotenv
load_dotenv()


# LLM loader
def load_llm(choice):
    if choice == "OpenAI (Paid)":
        return ChatOpenAI(
            temperature=0,
            streaming=True,
            api_key=os.getenv("OPENAI_API_KEY")
        )
    else:
        llm = HuggingFacePipeline.from_model_id(
        model_id='TinyLlama/TinyLlama-1.1B-Chat-v1.0',
        task='text-generation',
        pipeline_kwargs=dict(
            temperature=0.1,
            max_new_tokens=2048
            )
        )
        model = ChatHuggingFace(llm=llm)
        return model