Update app.py

app.py

@@ -1,31 +1,34 @@
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 app = FastAPI()
 app.add_middleware(
-    CORSMiddleware, allow_origins=["*"], allow_credentials=False,
+    CORSMiddleware,
+    allow_origins=["*"], allow_credentials=False,
     allow_methods=["*"], allow_headers=["*"]
 )
 
-model_name = "NeuralNovel/Mistral-7B-Instruct-v0.2-Neural-Story"
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name, trust_remote_code=True)
+model_name = "tiiuae/falcon-7b-instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
 class PromptRequest(BaseModel):
     prompt: str
 
 @app.post("/api/generate-story")
 async def generate_story(req: PromptRequest):
-    if not req.prompt.strip():
+    prompt = req.prompt.strip()
+    if not prompt:
         raise HTTPException(status_code=400, detail="Prompt must not be empty")
-    inputs = tokenizer(req.prompt, return_tensors="pt", truncation=True)
+
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
     outputs = model.generate(
         **inputs,
-        max_new_tokens=250,
-        temperature=0.9,
-        top_p=0.95,
+        max_new_tokens=200,
+        temperature=0.85,
+        top_p=0.9,
         repetition_penalty=1.2,
         do_sample=True
     )