Update app.py

app.py

@@ -2,51 +2,31 @@ from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import torch
 
 app = FastAPI()
-
-# Enable CORS for frontend fetch requests
 app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=False,
-    allow_methods=["*"],
-    allow_headers=["*"]
+  CORSMiddleware, allow_origins=["*"], allow_credentials=False,
+  allow_methods=["*"], allow_headers=["*"]
 )
 
-# Load FLAN-T5 model
-model_name = "google/flan-t5-base"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+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)
 
 class QueryRequest(BaseModel):
-    query: str
-    echo: bool = False
-
-@app.post("/api/query")
-async def generate_response(req: QueryRequest):
-    query = req.query.strip()
-    if not query:
-        raise HTTPException(status_code=400, detail="Query must not be empty")
-    
-    if req.echo:
-        return {"response": query}
-    
-    # Encode input
-    inputs = tokenizer(query, return_tensors="pt", truncation=True)
-
-    # Generate response with better decoding
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=150,
-        temperature=0.9,
-        top_p=0.95,
-        repetition_penalty=1.2,
-        do_sample=True,
-        num_return_sequences=1
-    )
+  prompt: str
 
-    # Decode output
-    generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"response": generated}
+@app.post("/api/generate-story")
+def generate_story(req: QueryRequest):
+  if not req.prompt.strip():
+    raise HTTPException(status_code=400, detail="Prompt must not be empty")
+  inputs = tokenizer(req.prompt, return_tensors="pt", truncation=True)
+  outputs = model.generate(
+    **inputs,
+    max_new_tokens=200,
+    temperature=0.9,
+    top_p=0.95,
+    repetition_penalty=1.2,
+    do_sample=True
+  )
+  return {"story": tokenizer.decode(outputs[0], skip_special_tokens=True)}