Final_Assignment_Template / model_factory.py
mjschock's picture
Enhance project structure by adding new files and updating dependencies. Introduce model_factory.py for streamlined model creation using environment variables. Update .gitignore to exclude additional cache and model files. Modify main_v2.py to utilize ModelFactory for model instantiation and adjust API key loading. Expand requirements.txt with new dependencies for improved functionality. Add new notebooks for fine-tuning and model training, enhancing project documentation and usability.
7461cd0 unverified
raw
history blame
1.24 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",
)