Update app.py

app.py

@@ -2,16 +2,29 @@ import os
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import (
+    AutoModelForCausalLM, 
+    AutoTokenizer,
+    Trainer,
+    TrainingArguments,
+    DataCollatorForLanguageModeling
+)
+from datasets import load_dataset
+from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training
 import torch
+from datetime import datetime
+import traceback
+
+# Environment setup
 os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
 os.environ["HF_HOME"] = "/app/cache"
 os.environ["XDG_CACHE_HOME"] = "/app/cache"
 os.makedirs("/app/cache", exist_ok=True)
+os.makedirs("/app/finetuned", exist_ok=True)
 
 app = FastAPI()
 
-# Enable CORS
+# CORS Configuration
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -19,37 +32,47 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Load model with caching
-model_name = "gpt2"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    device_map="auto",  # Automatically uses GPU if available
-    torch_dtype=torch.float16  # Optimize for GPU
-)
+# Model Loading with error handling
+try:
+    model_name = "gpt2"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+    
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        device_map="auto",
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+    )
+except Exception as e:
+    print(f"Model loading failed: {str(e)}")
+    raise
 
+# Pydantic Models
 class Query(BaseModel):
     message: str = Field(..., max_length=500)
 
+class FineTuneRequest(BaseModel):
+    epochs: int = Field(1, gt=0, le=5)
+    learning_rate: float = Field(5e-5, gt=0, le=1e-3)
+
+# Response Generation
 def generate_response(user_message):
-    prompt = f"User: {user_message}\nAI:"
-    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
+    prompt = f"<FIN_QA>Question: {user_message}\nAnswer:"
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     
-    output = model.generate(
-        input_ids,
-        max_new_tokens=100,
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=150,
         temperature=0.7,
         do_sample=True,
-        no_repeat_ngram_size=2,
+        no_repeat_ngram_size=3,
         repetition_penalty=1.5,
-        early_stopping=True,
         eos_token_id=tokenizer.eos_token_id
     )
     
-    full_response = tokenizer.decode(output[0], skip_special_tokens=True)
-    response = full_response.split("AI:")[-1].split("\nUser:")[0].strip()
-    return response or "I'm not sure how to respond to that."
+    return tokenizer.decode(outputs[0], skip_special_tokens=True).split("Answer:")[-1].strip()
 
+# API Endpoints
 @app.post("/chat")
 async def chat(query: Query):
     try:
@@ -58,6 +81,76 @@ async def chat(query: Query):
     except Exception as e:
         return {"error": str(e)}
 
+@app.post("/fine-tune")
+async def fine_tune(params: FineTuneRequest):
+    try:
+        # Load and combine datasets
+        alpaca = load_dataset("gbharti/finance-alpaca", split="train[:20%]")  # Sample 20% for demo
+        fiqa = load_dataset("bilalRahib/fiqa-personal-finance-dataset", "full", split="train[:20%]")
+        
+        # Formatting function
+        def format_example(ex):
+            if 'instruction' in ex:
+                return {"text": f"Instruction: {ex['instruction']}\nInput: {ex['input']}\nOutput: {ex['output']}"}
+            else:
+                return {"text": f"Question: {ex['question']}\nAnswer: {ex['answer']}"}
+        
+        dataset = alpaca.map(format_example) + fiqa.map(format_example)
+        
+        # Tokenize
+        def tokenize(ex):
+            return tokenizer(ex["text"], truncation=True, max_length=256, padding="max_length")
+        
+        dataset = dataset.map(tokenize, batched=True)
+        
+        # LoRA Configuration
+        peft_config = LoraConfig(
+            r=8,
+            lora_alpha=16,
+            target_modules=["c_attn", "c_proj", "c_fc"],
+            lora_dropout=0.05,
+            bias="none",
+            task_type="CAUSAL_LM"
+        )
+        
+        model = prepare_model_for_int8_training(model)
+        model = get_peft_model(model, peft_config)
+        
+        # Training
+        trainer = Trainer(
+            model=model,
+            args=TrainingArguments(
+                output_dir="/app/finetuned",
+                per_device_train_batch_size=2,
+                num_train_epochs=params.epochs,
+                learning_rate=params.learning_rate,
+                logging_dir="/app/logs",
+                save_strategy="epoch",
+                fp16=torch.cuda.is_available(),
+            ),
+            train_dataset=dataset,
+            data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
+        )
+        
+        trainer.train()
+        model.save_pretrained("/app/finetuned")
+        
+        return {
+            "status": "success",
+            "trained_samples": len(dataset),
+            "training_time": datetime.now().isoformat()
+        }
+        
+    except Exception as e:
+        return {
+            "error": str(e),
+            "traceback": traceback.format_exc()
+        }
+
 @app.get("/")
 def health_check():
-    return {"status": "OK"}
+    return {
+        "status": "healthy",
+        "model": model_name,
+        "device": str(model.device)
+    }