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