Update app.py

app.py

@@ -1,41 +1,89 @@
 from fastapi import FastAPI
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import os
+
+# ==========================================
+# 1. SETUP & MODEL LOADING (Executed once on startup)
+# ==========================================
 
 app = FastAPI()
 
-MODEL = "natalieparker/LumaAI-160M-v3"
+# CRITICAL: Since you merged the weights into the root folder, 
+# the model path inside the Hugging Face container is always '.'
+MODEL_DIR = "."
+
+# The model must be loaded into CPU memory first, then transferred if GPU is available.
+# We are forcing CPU inference as you intended.
+DEVICE = "cpu" 
+
+try:
+    print("🔄 Loading tokenizer...")
+    # Use MODEL_DIR ('.') because files are in the root of the deployed app
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, use_fast=True)
+    
+    print("🔄 Loading model on CPU...")
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_DIR,
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True # Optimization for CPU RAM
+    )
+    model.to(DEVICE)
+    print("✅ Model loaded successfully on CPU.")
+
+except Exception as e:
+    # If the model fails to load, print the error and set model to None
+    print(f"FATAL MODEL LOAD ERROR: {e}")
+    model = None
+    tokenizer = None
 
-print("🔄 Loading tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL, use_fast=True)
 
-print("🔄 Loading model on CPU...")
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL,
-    torch_dtype=torch.float32,   # CPU only
-    low_cpu_mem_usage=True
-)
-model.to("cpu")
+# ==========================================
+# 2. ENDPOINTS
+# ==========================================
 
 @app.get("/")
 def root():
-    return {"status": "LumaAI API is live on CPU"}
+    """Health check endpoint."""
+    return {"status": "LumaAI API is live", "model_loaded": model is not None}
 
 @app.post("/generate")
 def generate(prompt: str):
-    inputs = tokenizer(prompt, return_tensors="pt")
+    """Generates text response from the model."""
+    if model is None:
+        return {"error": "Model failed to load during startup."}
+        
+    # --- 1. PROMPT PREPARATION ---
+    # We use the final, clean format you trained on
+    formatted_prompt = f"User: {prompt}\nCharacter:"
+    
+    # --- 2. INFERENCE ---
+    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(DEVICE)
+    
     with torch.no_grad():
         output = model.generate(
             **inputs,
             max_new_tokens=150,
-            temperature=0.9,
+            temperature=0.75, # Balanced creativity
             top_p=0.9,
-            repetition_penalty=1.05,
-            do_sample=True
+            repetition_penalty=1.2,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
         )
+    
+    # --- 3. CLEANING ---
     text = tokenizer.decode(output[0], skip_special_tokens=True)
-    return {"response": text}
+    
+    # Extract response after the Character: tag
+    if "Character:" in text:
+        response = text.split("Character:")[-1]
+    else:
+        response = text
+        
+    # Clean up future user input and trim
+    response = response.split("User:")[0].strip()
+    
+    # Final punctuation polish (from our earlier fixes)
+    response = response.replace(" .", ".").replace(" ,", ",").replace(" ?", "?").replace(" !", "!")
 
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    return {"response": response}