Enhance serve.py to handle additional content types by converting dictionary text and joining list items. Update train.py to replace FastLanguageModel with FastModel and LiteLLMModel, streamline model loading, and adjust dataset preparation logic. Modify config.yaml to change max_samples for testing and add provider information for model configuration.

conf/config.yaml

@@ -4,7 +4,9 @@ defaults:
 # Model configuration
 model:
   name: "unsloth/SmolLM2-135M-Instruct-bnb-4bit"
+  # name: "HuggingFaceTB/SmolLM2-135M-Instruct"
   max_seq_length: 2048  # Auto supports RoPE Scaling internally
+  provider: "openai"
   dtype: null  # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
   load_in_4bit: true  # Use 4bit quantization to reduce memory usage
 
@@ -77,5 +79,5 @@ test_dataset:
   name: "gaia-benchmark/GAIA"
   config: "2023_level1"  # Use level 1 questions for testing
   split: "test"  # Use test split for testing
-  max_samples: 10  # Number of samples to test on
+  max_samples: 3  # Number of samples to test on
   max_length: 2048  # Maximum sequence length for testing

serve.py

@@ -153,6 +153,12 @@ class ModelDeployment:
 
                     content["type"] = "image"
                     del content["image_url"]
+                elif isinstance(content, dict) and "text" in content:
+                    # Convert content to string if it's a dict with text
+                    message["content"] = content["text"]
+                elif isinstance(content, list):
+                    # Join list items with newlines if content is a list
+                    message["content"] = "\n".join(content)
 
         images = images if images else None
 

train.py

@@ -23,7 +23,7 @@ import hydra
 from omegaconf import DictConfig, OmegaConf
 
 # isort: off
-from unsloth import FastLanguageModel, is_bfloat16_supported  # noqa: E402
+from unsloth import FastLanguageModel, FastModel, is_bfloat16_supported  # noqa: E402
 from unsloth.chat_templates import get_chat_template  # noqa: E402
 
 # isort: on
@@ -39,7 +39,7 @@ from datasets import (
     load_dataset,
 )
 from peft import PeftModel
-from smolagents import CodeAgent, Model, TransformersModel, VLLMModel
+from smolagents import CodeAgent, LiteLLMModel, Model, TransformersModel, VLLMModel
 from smolagents.monitoring import LogLevel
 from transformers import (
     AutoModelForCausalLM,
@@ -97,7 +97,7 @@ def load_model(cfg: DictConfig) -> tuple[FastLanguageModel, AutoTokenizer]:
     """Load and configure the model."""
     logger.info("Loading model and tokenizer...")
     try:
-        model, tokenizer = FastLanguageModel.from_pretrained(
+        model, tokenizer = FastModel.from_pretrained(
             model_name=cfg.model.name,
             max_seq_length=cfg.model.max_seq_length,
             dtype=cfg.model.dtype,
@@ -106,7 +106,7 @@ def load_model(cfg: DictConfig) -> tuple[FastLanguageModel, AutoTokenizer]:
         logger.info("Base model loaded successfully")
 
         # Configure LoRA
-        model = FastLanguageModel.get_peft_model(
+        model = FastModel.get_peft_model(
             model,
             r=cfg.peft.r,
             target_modules=cfg.peft.target_modules,
@@ -242,19 +242,19 @@ def main(cfg: DictConfig) -> None:
         logger.info(f"Configuration:\n{OmegaConf.to_yaml(cfg)}")
 
         # Install dependencies
-        install_dependencies()
+        # install_dependencies()
 
-        # Load model and tokenizer
-        model, tokenizer = load_model(cfg)
+        # Train if requested
+        if cfg.train:
+            # Load model and tokenizer
+            model, tokenizer = load_model(cfg)
 
-        # Load and prepare dataset
-        dataset, tokenizer = load_and_format_dataset(tokenizer, cfg)
+            # Load and prepare dataset
+            dataset, tokenizer = load_and_format_dataset(tokenizer, cfg)
 
-        # Create trainer
-        trainer: Trainer = create_trainer(model, tokenizer, dataset, cfg)
+            # Create trainer
+            trainer: Trainer = create_trainer(model, tokenizer, dataset, cfg)
 
-        # Train if requested
-        if cfg.train:
             logger.info("Starting training...")
             trainer.train()
 
@@ -304,8 +304,11 @@ def main(cfg: DictConfig) -> None:
                 torch.cuda.empty_cache()
 
                 # Initialize model
-                model: Model = Model(
-                    model_id=cfg.model.name,
+                model: Model = LiteLLMModel(
+                    api_base="http://localhost:8000/v1",
+                    api_key="not-needed",
+                    model_id=f"{cfg.model.provider}/{cfg.model.name}",
+                    # model_id=cfg.model.name,
                     # model_id=cfg.output.dir,
                 )