File size: 674 Bytes
95be894 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import pipeline
# Initialize FastAPI app
app = FastAPI()
# Pydantic schema for input data
class Query(BaseModel):
prompt: str
# Load a Hugging Face pipeline
model = pipeline("text-generation", model="gpt2")
# Define a POST endpoint for processing queries
@app.post("/generate/")
async def generate_text(query: Query):
try:
response = model(query.prompt, max_length=150, num_return_sequences=1)
return {"response": response[0]["generated_text"]}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
|