Update app.py

app.py

@@ -4,7 +4,7 @@ from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import os
 
-# Use a writable folder for offloading weights (Hugging Face Spaces restricts /app)
+# Use a writable folder for offloading weights
 offload_dir = "/tmp/offload"
 os.makedirs(offload_dir, exist_ok=True)
 
@@ -13,28 +13,24 @@ app = FastAPI()
 # CORS setup
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Allow all origins for testing
+    allow_origins=["*"],
     allow_credentials=False,
     allow_methods=["*"],
     allow_headers=["*"]
 )
 
-# Model name (7B model - large, will offload to /tmp)
-model_name = "ethzanalytics/RedPajama-INCITE-7B-Base-sharded-bf16"
+# Smaller & faster model
+model_name = "tiiuae/falcon-rw-1b"
 
-# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-# Load model with /tmp offload folder
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    torch_dtype="bfloat16",
+    torch_dtype="auto",
     device_map="auto",
     low_cpu_mem_usage=True,
     offload_folder=offload_dir
 )
 
-# Request body schema
 class PromptRequest(BaseModel):
     prompt: str
 
@@ -44,19 +40,13 @@ async def generate_story(req: PromptRequest):
     if not prompt:
         raise HTTPException(status_code=400, detail="Prompt must not be empty")
 
-    # Tokenize input
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True).to(model.device)
-
-    # Generate story
     outputs = model.generate(
         **inputs,
-        max_new_tokens=200,
+        max_new_tokens=150,
         do_sample=True,
         temperature=0.9,
-        top_p=0.9,
-        repetition_penalty=1.2
+        top_p=0.9
     )
-
-    # Decode and return
     story = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return {"story": story}