Update main.py

main.py

@@ -4,7 +4,7 @@ import uvicorn
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import time
 import uuid
 import logging
@@ -12,7 +12,6 @@ import logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Request/Response Models
 class GenerateRequest(BaseModel):
     prompt: str
     max_tokens: int = 512
@@ -24,14 +23,12 @@ class GenerateResponse(BaseModel):
     text: str
     created: int
 
-# Model Manager
 class ModelManager:
     def __init__(self):
         self.model = None
         self.tokenizer = None
         self.device = None
         self.loading_status = "not_started"
-        # Use a model from Hugging Face Hub
         self.model_id = os.environ.get("MODEL_ID", "ugonna/llama3.18B-Fine-tunedByUgo3")
         
     def setup_device(self):
@@ -40,7 +37,7 @@ class ModelManager:
             logger.info(f"✅ Using GPU: {torch.cuda.get_device_name(0)}")
         else:
             self.device = torch.device("cpu")
-            logger.info("⚠️ Using CPU (slower for LLMs)")
+            logger.info("⚠️ Using CPU - this will be VERY slow for large models")
     
     def load_model(self):
         try:
@@ -52,20 +49,27 @@ class ModelManager:
             if self.tokenizer.pad_token is None:
                 self.tokenizer.pad_token = self.tokenizer.eos_token
             
-            logger.info(f"📥 Loading model from {self.model_id}...")
+            logger.info(f"📥 Loading model from {self.model_id} with quantization...")
+            
+            # Use 8-bit quantization to reduce memory (requires bitsandbytes)
+            bnb_config = BitsAndBytesConfig(
+                load_in_8bit=True,
+                llm_int8_threshold=6.0
+            )
             
-            # Load with optimizations for CPU
             self.model = AutoModelForCausalLM.from_pretrained(
                 self.model_id,
-                torch_dtype=torch.float32,  # Use float32 for CPU
+                quantization_config=bnb_config if torch.cuda.is_available() else None,
+                device_map="auto" if torch.cuda.is_available() else None,
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                 low_cpu_mem_usage=True,
                 trust_remote_code=True
             )
             
-            # Move to CPU explicitly
-            self.model = self.model.to(self.device)
-            self.model.eval()
+            if not torch.cuda.is_available():
+                self.model = self.model.to(self.device)
             
+            self.model.eval()
             self.loading_status = "loaded"
             logger.info("✅ Model loaded successfully!")
             return True
@@ -73,6 +77,7 @@ class ModelManager:
         except Exception as e:
             self.loading_status = "failed"
             logger.error(f"❌ Failed to load model: {e}")
+            logger.error("This model may be too large for the free tier. Consider using a smaller model.")
             return False
     
     def generate(self, prompt: str, max_tokens: int = 512, temperature: float = 0.7, top_p: float = 0.9):
@@ -80,15 +85,17 @@ class ModelManager:
             raise Exception("Model not loaded")
         
         start_time = time.time()
-        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
+        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
         
-        # Move inputs to the same device as model
-        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        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()}
         
         with torch.no_grad():
             outputs = self.model.generate(
                 **inputs,
-                max_new_tokens=max_tokens,
+                max_new_tokens=min(max_tokens, 200),  # Limit for CPU
                 temperature=temperature,
                 top_p=top_p,
                 do_sample=True,
@@ -104,13 +111,10 @@ class ModelManager:
         
         return generated_text
 
-# Initialize model manager
 model_manager = ModelManager()
 
-# FastAPI app
 app = FastAPI(title="NAI Bot API", version="1.0.0")
 
-# CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -121,7 +125,6 @@ app.add_middleware(
 
 @app.on_event("startup")
 async def startup_event():
-    """Load model in background on startup"""
     import threading
     thread = threading.Thread(target=model_manager.load_model, daemon=True)
     thread.start()
@@ -132,13 +135,14 @@ async def root():
         "message": "NAI Bot API is running",
         "docs": "/docs",
         "status": model_manager.loading_status,
-        "model": model_manager.model_id
+        "model": model_manager.model_id,
+        "warning": "Large model on free tier may fail due to memory limits"
     }
 
 @app.get("/health")
 async def health_check():
     return {
-        "status": "healthy" if model_manager.model is not None else "loading",
+        "status": "healthy" if model_manager.model is not None else "failed",
         "model_loaded": model_manager.model is not None,
         "loading_status": model_manager.loading_status,
         "device": str(model_manager.device) if model_manager.device else "unknown",
@@ -148,7 +152,10 @@ async def health_check():
 @app.post("/generate", response_model=GenerateResponse)
 async def generate(request: GenerateRequest):
     if model_manager.model is None:
-        raise HTTPException(status_code=503, detail=f"Model is still loading (status: {model_manager.loading_status}). Try again in a few seconds.")
+        raise HTTPException(
+            status_code=503, 
+            detail=f"Model failed to load. Your 8B model is too large for free tier. Please switch to a smaller model like 'microsoft/DialoGPT-small'"
+        )
     
     if not request.prompt:
         raise HTTPException(status_code=400, detail="Prompt cannot be empty")