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