Update app.py

app.py

@@ -1,74 +1,78 @@
 import os
 import torch
-from flask import Flask, request, jsonify
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List, Dict, Optional
 from datasets import load_dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import uvicorn
 
-app = Flask(__name__)
+app = FastAPI()
 
-# Global variables to store model, tokenizer, and dataset
+# Global variables
 model = None
 tokenizer = None
 dataset = None
 
-# Function to load the model and dataset
-def load_model_and_data():
-    global model, tokenizer, dataset
-    
-    # Load the base model and tokenizer
-    model_id = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    
-    # Load the model with reduced precision for efficiency
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16,
-        device_map="auto"
-    )
-    
-    # Load your dataset
-    dataset = load_dataset("lahiruchamika27/tia")
-    print("Model, tokenizer, and dataset loaded successfully!")
+# Pydantic models for request/response
+class ChatTurn(BaseModel):
+    user: Optional[str] = None
+    assistant: Optional[str] = None
 
-# Initialize the model on startup
-@app.before_request
-def before_request():
-    global model, tokenizer, dataset
-    if model is None or tokenizer is None or dataset is None:
-        load_model_and_data()
+class ChatRequest(BaseModel):
+    message: str
+    history: Optional[List[ChatTurn]] = []
 
-# Define chat endpoint
-@app.route('/api/chat', methods=['POST'])
-def chat():
+class ChatResponse(BaseModel):
+    response: str
+
+# Load model on startup
+@app.on_event("startup")
+async def startup_event():
+    global model, tokenizer, dataset
     try:
-        # Get input from request
-        data = request.json
-        if not data or 'message' not in data:
-            return jsonify({"error": "No message provided"}), 400
+        # Load the model and tokenizer
+        model_id = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
         
-        user_message = data['message']
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
         
-        # You can optionally retrieve conversation history
-        conversation_history = data.get('history', [])
-        
-        # Create input for the model
-        # Format may need adjustment based on your model's expected format
-        if conversation_history:
+        # Load dataset
+        dataset = load_dataset("lahiruchamika27/tia")
+        print("Model, tokenizer, and dataset loaded successfully!")
+    except Exception as e:
+        print(f"Error loading model: {str(e)}")
+        # Continue without failing - we'll handle errors in the endpoints
+
+@app.post("/api/chat", response_model=ChatResponse)
+async def chat(request: ChatRequest):
+    global model, tokenizer
+    
+    # Ensure model is loaded
+    if model is None or tokenizer is None:
+        raise HTTPException(status_code=500, detail="Model or tokenizer not loaded")
+    
+    try:
+        # Format conversation
+        if request.history:
             full_prompt = ""
-            for turn in conversation_history:
-                if 'user' in turn:
-                    full_prompt += f"User: {turn['user']}\n"
-                if 'assistant' in turn:
-                    full_prompt += f"Assistant: {turn['assistant']}\n"
+            for turn in request.history:
+                if turn.user:
+                    full_prompt += f"User: {turn.user}\n"
+                if turn.assistant:
+                    full_prompt += f"Assistant: {turn.assistant}\n"
             
-            full_prompt += f"User: {user_message}\nAssistant:"
+            full_prompt += f"User: {request.message}\nAssistant:"
         else:
-            full_prompt = f"User: {user_message}\nAssistant:"
+            full_prompt = f"User: {request.message}\nAssistant:"
         
         # Tokenize and generate
         inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
         
-        # Generate response
         with torch.no_grad():
             outputs = model.generate(
                 inputs["input_ids"],
@@ -81,29 +85,33 @@ def chat():
         # Decode the output
         response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
         
-        return jsonify({"response": response.strip()})
+        return ChatResponse(response=response.strip())
     
     except Exception as e:
-        return jsonify({"error": str(e)}), 500
+        raise HTTPException(status_code=500, detail=str(e))
 
-# Optional endpoint to get examples from your dataset
-@app.route('/api/examples', methods=['GET'])
-def get_examples():
+@app.get("/api/examples")
+async def get_examples(count: int = 5, split: str = "train"):
+    global dataset
+    
+    if dataset is None:
+        raise HTTPException(status_code=500, detail="Dataset not loaded")
+    
     try:
-        # Get a sample from the dataset
-        num_examples = int(request.args.get('count', 5))
-        split = request.args.get('split', 'train')
-        
         if split in dataset:
-            examples = dataset[split][:num_examples]
-            return jsonify({"examples": examples})
+            # Convert dataset items to dict for easier JSON serialization
+            examples = [dict(item) for item in dataset[split][:count]]
+            return {"examples": examples}
         else:
-            return jsonify({"error": f"Split '{split}' not found in dataset"}), 400
+            raise HTTPException(status_code=400, detail=f"Split '{split}' not found in dataset")
     
     except Exception as e:
-        return jsonify({"error": str(e)}), 500
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/health")
+async def health_check():
+    return {"status": "ok", "model_loaded": model is not None, "tokenizer_loaded": tokenizer is not None}
 
-if __name__ == '__main__':
-    # Get port from environment variable for HF Spaces compatibility
-    port = int(os.environ.get('PORT', 7860))
-    app.run(host='0.0.0.0', port=port)
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 7860))
+    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=False)