Final_Assignment_Template / model_factory.py
mjschock's picture
Refactor imports in main_v2.py to include SmartSearchTool, enhance code organization in model_factory.py by removing unnecessary whitespace, and improve formatting in test_questions.py for better readability and consistency.
352f525 unverified
raw
history blame contribute delete
1.22 kB
import os
from dotenv import find_dotenv, load_dotenv
from smolagents import LiteLLMModel, TransformersModel
class ModelFactory:
@staticmethod
def create_model():
"""
Creates and returns a LiteLLMModel instance configured with environment variables.
Returns:
LiteLLMModel: A configured instance of LiteLLMModel
"""
# Load environment variables
load_dotenv(find_dotenv())
# Get configuration from environment variables
api_base = os.getenv("API_BASE")
api_key = os.getenv("API_KEY")
model_id = os.getenv("MODEL_ID")
# Create and return the model
# return LiteLLMModel(
# api_base=api_base,
# api_key=api_key,
# model_id=model_id,
# )
return TransformersModel(
# max_new_tokens=5000,
max_new_tokens=256,
model_id="HuggingFaceTB/SmolLM2-135M-Instruct",
# model_id="HuggingFaceTB/SmolLM2-360M-Instruct",
# model_id="HuggingFaceTB/SmolLM2-1.7B-Instruct",
# model_id="HuggingFaceTB/SmolVLM2-256M-Video-Instruct",
# model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
)