app.py

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

app = FastAPI()
app.add_middleware(
  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)

class QueryRequest(BaseModel):
  prompt: str

@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)}