from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from model import parth

app = FastAPI(
    title="Text Generation API",
    description="API for generating text using DistilGPT-2 model",
    version="1.0.0"
)

class GenerateRequest(BaseModel):
    prompt: str
    max_length: int = 50

class GenerateResponse(BaseModel):
    generated_text: str

@app.post("/generate", response_model=GenerateResponse)
async def generate(request: GenerateRequest):
    try:
        generated_text = parth.generate_text(request.prompt, request.max_length)
        return GenerateResponse(generated_text=generated_text)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)