Update main.py

main.py

@@ -3,84 +3,94 @@ import os
 import uvicorn
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel, Field
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import time
 import uuid
-from typing import Dict, Any
 import logging
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Request/Response Models
+class GenerateRequest(BaseModel):
+    prompt: str
+    max_tokens: int = 512
+    temperature: float = 0.7
+    top_p: float = 0.9
+
+class GenerateResponse(BaseModel):
+    id: str
+    text: str
+    created: int
+
 # Model Manager
 class ModelManager:
     def __init__(self):
         self.model = None
         self.tokenizer = None
         self.device = None
-        self.model_path = "/opt/render/project/src/models"
         self.loading_status = "not_started"
-    
+        # Use a model from Hugging Face Hub
+        self.model_id = os.environ.get("MODEL_ID", "microsoft/DialoGPT-small")
+        
     def setup_device(self):
         if torch.cuda.is_available():
             self.device = torch.device("cuda")
-            logger.info(f"Using GPU: {torch.cuda.get_device_name(0)}")
+            logger.info(f"✅ Using GPU: {torch.cuda.get_device_name(0)}")
         else:
             self.device = torch.device("cpu")
-            logger.info("Using CPU")
+            logger.info("⚠️ Using CPU (slower for LLMs)")
     
     def load_model(self):
         try:
             self.loading_status = "loading"
             self.setup_device()
             
-            logger.info("Loading tokenizer...")
-            self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
+            logger.info(f"📥 Loading tokenizer from {self.model_id}...")
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
             if self.tokenizer.pad_token is None:
                 self.tokenizer.pad_token = self.tokenizer.eos_token
             
-            logger.info("Loading model...")
-            if torch.cuda.is_available():
-                self.model = AutoModelForCausalLM.from_pretrained(
-                    self.model_path,
-                    torch_dtype=torch.float16,
-                    device_map="auto",
-                    trust_remote_code=True
-                )
-            else:
-                self.model = AutoModelForCausalLM.from_pretrained(
-                    self.model_path,
-                    torch_dtype=torch.float32,
-                    trust_remote_code=True
-                ).to(self.device)
+            logger.info(f"📥 Loading model from {self.model_id}...")
+            
+            # Load with optimizations for CPU
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_id,
+                torch_dtype=torch.float32,  # Use float32 for CPU
+                low_cpu_mem_usage=True,
+                trust_remote_code=True
+            )
+            
+            # Move to CPU explicitly
+            self.model = self.model.to(self.device)
+            self.model.eval()
             
             self.loading_status = "loaded"
             logger.info("✅ Model loaded successfully!")
             return True
+            
         except Exception as e:
             self.loading_status = "failed"
-            logger.error(f"Failed to load model: {e}")
+            logger.error(f"❌ Failed to load model: {e}")
             return False
     
-    def generate(self, prompt: str, **kwargs):
+    def generate(self, prompt: str, max_tokens: int = 512, temperature: float = 0.7, top_p: float = 0.9):
         if self.model is None:
             raise Exception("Model not loaded")
         
         start_time = time.time()
-        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
+        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
         
-        if torch.cuda.is_available():
-            inputs = {k: v.cuda() for k, v in inputs.items()}
-        else:
-            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        # Move inputs to the same device as model
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
         
         with torch.no_grad():
             outputs = self.model.generate(
                 **inputs,
-                max_new_tokens=kwargs.get('max_tokens', 512),
-                temperature=kwargs.get('temperature', 0.7),
-                top_p=kwargs.get('top_p', 0.9),
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
                 do_sample=True,
                 pad_token_id=self.tokenizer.pad_token_id,
                 eos_token_id=self.tokenizer.eos_token_id
@@ -89,14 +99,18 @@ class ModelManager:
         generated_ids = outputs[0][inputs['input_ids'].shape[1]:]
         generated_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
         
-        return {"text": generated_text}
+        elapsed = time.time() - start_time
+        logger.info(f"✅ Generated {len(generated_ids)} tokens in {elapsed:.2f}s")
+        
+        return generated_text
 
 # Initialize model manager
 model_manager = ModelManager()
 
-# Create FastAPI app
-app = FastAPI(title="Model API", version="1.0.0")
+# FastAPI app
+app = FastAPI(title="NAI Bot API", version="1.0.0")
 
+# CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -107,42 +121,55 @@ app.add_middleware(
 
 @app.on_event("startup")
 async def startup_event():
-    """Load model on startup"""
+    """Load model in background on startup"""
     import threading
     thread = threading.Thread(target=model_manager.load_model, daemon=True)
     thread.start()
 
+@app.get("/")
+async def root():
+    return {
+        "message": "NAI Bot API is running",
+        "docs": "/docs",
+        "status": model_manager.loading_status,
+        "model": model_manager.model_id
+    }
+
 @app.get("/health")
 async def health_check():
     return {
         "status": "healthy" if model_manager.model is not None else "loading",
         "model_loaded": model_manager.model is not None,
-        "device": str(model_manager.device) if model_manager.device else "unknown"
+        "loading_status": model_manager.loading_status,
+        "device": str(model_manager.device) if model_manager.device else "unknown",
+        "model_id": model_manager.model_id
     }
 
-@app.post("/generate")
-async def generate(request: dict):
+@app.post("/generate", response_model=GenerateResponse)
+async def generate(request: GenerateRequest):
     if model_manager.model is None:
-        raise HTTPException(status_code=503, detail="Model not ready")
+        raise HTTPException(status_code=503, detail=f"Model is still loading (status: {model_manager.loading_status}). Try again in a few seconds.")
+    
+    if not request.prompt:
+        raise HTTPException(status_code=400, detail="Prompt cannot be empty")
     
     try:
-        result = model_manager.generate(
-            prompt=request.get("prompt", ""),
-            max_tokens=request.get("max_tokens", 512),
-            temperature=request.get("temperature", 0.7)
+        generated_text = model_manager.generate(
+            prompt=request.prompt,
+            max_tokens=request.max_tokens,
+            temperature=request.temperature,
+            top_p=request.top_p
+        )
+        
+        return GenerateResponse(
+            id=str(uuid.uuid4()),
+            text=generated_text,
+            created=int(time.time())
         )
-        return {
-            "id": str(uuid.uuid4()),
-            "text": result["text"],
-            "created": int(time.time())
-        }
     except Exception as e:
+        logger.error(f"Generation error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
-@app.get("/")
-async def root():
-    return {"message": "Model API is running", "docs": "/docs"}
-
 if __name__ == "__main__":
-    port = int(os.environ.get("PORT", 8000))
-    uvicorn.run(app, host="0.0.0.0", port=port)
+    port = int(os.environ.get("PORT", 7860))
+    uvicorn.run(app, host="0.0.0.0", port=port)