Update app.py

app.py

@@ -1,8 +1,23 @@
+# Install required dependencies
+import subprocess
+import sys
+
+def install_packages():
+    packages = ["sentencepiece", "protobuf", "transformers", "torch", "accelerate"]
+    for package in packages:
+        try:
+            __import__(package)
+        except ImportError:
+            print(f"Installing {package}...")
+            subprocess.check_call([sys.executable, "-m", "pip", "install", package])
+
+install_packages()
+
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import uvicorn
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import time
 from fastapi.middleware.cors import CORSMiddleware
 
@@ -24,8 +39,7 @@ app.add_middleware(
 
 class YAHBot:
     def __init__(self):
-        # ✅ Changed to load from your HF repo instead of direct model name
-        self.repo_id = "Adedoyinjames/brain-ai"  # Your HF repo
+        self.repo_id = "Adedoyinjames/brain-ai"
         self.tokenizer = None
         self.model = None
         self._load_model()
@@ -34,8 +48,16 @@ class YAHBot:
         """Load the model from your Hugging Face repo"""
         try:
             print(f"🔄 Loading AI model from {self.repo_id}...")
-            self.tokenizer = AutoTokenizer.from_pretrained(self.repo_id)
-            self.model = AutoModelForSeq2SeqLM.from_pretrained(self.repo_id)
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.repo_id,
+                trust_remote_code=True  # Required for phi-3
+            )
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.repo_id,
+                trust_remote_code=True,  # Required for phi-3
+                torch_dtype=torch.float16,
+                device_map="auto"
+            )
             print("✅ AI model loaded successfully from HF repo!")
         except Exception as e:
             print(f"❌ Failed to load AI model from repo: {e}")
@@ -43,12 +65,12 @@ class YAHBot:
             self.tokenizer = None
     
     def generate_response(self, user_input):
-        """Generate response using AI model"""
+        """Generate response using causal language model"""
         if self.model and self.tokenizer:
             try:
-                prompt = f"Question: {user_input}\nAnswer: "
+                # Format prompt for phi-3 (causal LM)
+                prompt = f"<|user|>\n{user_input}<|end|>\n<|assistant|>\n"
                 
-                # Tokenize
                 inputs = self.tokenizer(
                     prompt, 
                     return_tensors="pt", 
@@ -57,18 +79,27 @@ class YAHBot:
                     padding=True
                 )
                 
-                # Generate response
+                # Move to same device as model
+                device = next(self.model.parameters()).device
+                inputs = {k: v.to(device) for k, v in inputs.items()}
+                
                 with torch.no_grad():
                     outputs = self.model.generate(
                         inputs.input_ids,
-                        max_length=150,
+                        max_new_tokens=150,
                         num_return_sequences=1,
                         temperature=0.7,
                         do_sample=True,
-                        pad_token_id=self.tokenizer.pad_token_id,
+                        pad_token_id=self.tokenizer.eos_token_id,  # Use EOS token for padding
+                        eos_token_id=self.tokenizer.eos_token_id,
                     )
                 
                 response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+                
+                # Remove the prompt from the response for cleaner output
+                if prompt in response:
+                    response = response.replace(prompt, "").strip()
+                
                 return response
                 
             except Exception as e:
@@ -100,7 +131,8 @@ async def root():
     return {
         "message": "YAH Tech AI API is running",
         "status": "active",
-        "model_repo": yah_bot.repo_id,  # Show which repo is being used
+        "model_repo": yah_bot.repo_id,
+        "model_type": "causal_lm",
         "endpoints": {
             "chat": "POST /api/chat",
             "health": "GET /api/health"