File size: 838 Bytes
9f0bffc |
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 |
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)
|