Update app.py

app.py

@@ -3,26 +3,25 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 
-# Configuration - Using even smaller model
-MODEL_ID = "google/gemma-1.1-2b-it"
+# Configuration
+MODEL_ID = "google/gemma-1.1-2b-it"  # Using smaller 2B version
 HF_TOKEN = os.getenv("HF_TOKEN")
-MAX_TOKENS = 80  # Very conservative limit
+MAX_TOKENS = 80  # Conservative limit
 
 def load_model():
-    """Simplified model loading that works with Spaces"""
-    print("🔄 Loading model (this may take a few minutes)...")
+    """Simplified model loading that works in Spaces"""
+    print("🔄 Loading model...")
     
-    # Load with explicit CPU mapping
+    # Load tokenizer
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
+    
+    # Explicit CPU-only loading
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
-        device_map="cpu",  # Force CPU-only
         torch_dtype=torch.float32,  # Required for CPU
         token=HF_TOKEN
-    )
+    ).to('cpu')  # Explicit CPU placement
     
-    # Ensure weights are tied
-    model.tie_weights()
     print("✅ Model loaded successfully!")
     return tokenizer, model
 
@@ -31,14 +30,13 @@ tokenizer, model = load_model()
 def predict(topic):
     """Memory-safe generation"""
     try:
-        prompt = f"Create a VERY short script about {topic}:\n1) Hook\n2) Point\n3) CTA\n\nScript:"
-        inputs = tokenizer(prompt, return_tensors="pt")
+        prompt = f"Create a short script about {topic}:\n1) Hook\n2) Point\n3) CTA\n\nScript:"
+        inputs = tokenizer(prompt, return_tensors="pt").to('cpu')
         
         outputs = model.generate(
             **inputs,
             max_new_tokens=MAX_TOKENS,
-            temperature=0.7,
-            do_sample=True
+            temperature=0.7
         )
         
         return tokenizer.decode(outputs[0], skip_special_tokens=True)