Update requirements.txt to include sse-starlette dependency, enhance serve.py with additional imports for FastLanguageModel and FastVisionModel, and refactor train.py for improved organization and memory tracking during model training.

requirements.txt

@@ -28,6 +28,7 @@ python-dotenv>=1.0.0
 requests>=2.32.3
 sentence-transformers>=4.1.0
 smolagents[litellm,telemetry,vllm]>=1.14.0
+sse-starlette>=2.3.4
 tensorboardX>=2.6.2.2
 trl>=0.17.0
 typing-extensions>=4.5.0

serve.py

@@ -6,6 +6,16 @@ from pprint import pprint
 from threading import Thread
 from typing import Any, Dict, List
 
+# isort: off
+from unsloth import (
+    FastLanguageModel,
+    FastVisionModel,
+    is_bfloat16_supported,
+)  # noqa: E402
+from unsloth.chat_templates import get_chat_template  # noqa: E402
+
+# isort: on
+
 from fastapi import FastAPI, Request
 from openai.types.chat.chat_completion import ChatCompletion
 from openai.types.chat.chat_completion import Choice as ChatCompletionChoice
@@ -20,7 +30,6 @@ from sse_starlette import EventSourceResponse
 from starlette.responses import JSONResponse
 from transformers.generation.streamers import AsyncTextIteratorStreamer
 from transformers.image_utils import load_image
-from unsloth import FastVisionModel
 
 dtype = (
     None  # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+

train.py

@@ -28,6 +28,9 @@ from unsloth.chat_templates import get_chat_template  # noqa: E402
 
 # isort: on
 
+import os
+
+import torch
 from datasets import (
     Dataset,
     DatasetDict,
@@ -35,20 +38,19 @@ from datasets import (
     IterableDatasetDict,
     load_dataset,
 )
+from peft import PeftModel
+from smolagents import CodeAgent, Model, TransformersModel, VLLMModel
+from smolagents.monitoring import LogLevel
 from transformers import (
+    AutoModelForCausalLM,
     AutoTokenizer,
     DataCollatorForLanguageModeling,
     Trainer,
     TrainingArguments,
-    AutoModelForCausalLM,
 )
 from trl import SFTTrainer
-from peft import PeftModel
-from smolagents import CodeAgent, Model, TransformersModel, VLLMModel
+
 from tools.smart_search.tool import SmartSearchTool
-from smolagents.monitoring import LogLevel
-import torch
-import os
 
 
 # Setup logging
@@ -259,13 +261,11 @@ def main(cfg: DictConfig) -> None:
             # Save model
             logger.info(f"Saving final model to {cfg.output.dir}...")
             trainer.save_model(cfg.output.dir)
-            
+
             # Save model in VLLM format
             logger.info("Saving model in VLLM format...")
             model.save_pretrained_merged(
-                cfg.output.dir,
-                tokenizer,
-                save_method="merged_16bit"
+                cfg.output.dir, tokenizer, save_method="merged_16bit"
             )
 
             # Print final metrics
@@ -284,10 +284,12 @@ def main(cfg: DictConfig) -> None:
             try:
                 # Enable memory history tracking
                 torch.cuda.memory._record_memory_history()
-                
+
                 # Set memory allocation configuration
-                os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True,max_split_size_mb:128'
-                
+                os.environ["PYTORCH_CUDA_ALLOC_CONF"] = (
+                    "expandable_segments:True,max_split_size_mb:128"
+                )
+
                 # Load test dataset
                 test_dataset = load_dataset(
                     cfg.test_dataset.name,
@@ -358,10 +360,10 @@ Please format your response as a JSON object with two keys:
                     try:
                         # Clear CUDA cache before each sample
                         torch.cuda.empty_cache()
-                        
+
                         # Format the task
-                        task = format_task(example['Question'])
-                        
+                        task = format_task(example["Question"])
+
                         # Run the agent
                         result = agent.run(
                             task=task,
@@ -372,20 +374,25 @@ Please format your response as a JSON object with two keys:
 
                         # Parse the result
                         import json
-                        json_str = result[result.find("{"):result.rfind("}")+1]
+
+                        json_str = result[result.find("{") : result.rfind("}") + 1]
                         parsed_result = json.loads(json_str)
                         answer = parsed_result["succinct_answer"]
-                        
+
                         logger.info(f"\nTest Sample {i+1}:")
                         logger.info(f"Question: {example['Question']}")
                         logger.info(f"Model Response: {answer}")
                         logger.info("-" * 80)
-                        
+
                         # Log memory usage after each sample
                         logger.info(f"Memory usage after sample {i+1}:")
-                        logger.info(f"Allocated: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
-                        logger.info(f"Reserved: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")
-                            
+                        logger.info(
+                            f"Allocated: {torch.cuda.memory_allocated() / 1024**2:.2f} MB"
+                        )
+                        logger.info(
+                            f"Reserved: {torch.cuda.memory_reserved() / 1024**2:.2f} MB"
+                        )
+
                     except Exception as e:
                         logger.error(f"Error processing test sample {i+1}: {str(e)}")
                         continue