Spaces:
Sleeping
Sleeping
File size: 1,380 Bytes
d86fc01 1c18d40 d86fc01 52808f5 d86fc01 1c18d40 d86fc01 52808f5 1c18d40 d86fc01 1c18d40 d86fc01 1c18d40 d86fc01 1c18d40 52808f5 1c18d40 52808f5 1c18d40 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
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=["*"]
)
# Load FLAN-T5 model
model_name = "google/flan-t5-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
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
)
# Decode output
generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
return {"response": generated}
|