Spaces:
Sleeping
Sleeping
File size: 802 Bytes
80e8504 |
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 |
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import pipeline
app = FastAPI()
# Initialize the text-generation pipeline
# We can load this on startup to avoid loading it on every request,
# although pipeline does caching, doing it globally is standard for this simple app.
generator = pipeline("text-generation", model="distilgpt2")
class ChatRequest(BaseModel):
text: str
@app.get("/")
def read_root():
return {"status": "Cool Shot Systems Backend Online"}
@app.post("/chat")
def chat(request: ChatRequest):
try:
# Generate response
response = generator(request.text, max_length=100, num_return_sequences=1)
return response[0]
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
|