1MR's picture
Create app.py
2b56f22 verified
raw
history blame
944 Bytes
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from transformers import pipeline
from fastapi.requests import Request
app = FastAPI()
# Initialize the text generation pipeline
pipe = pipeline("text2text-generation", model="google/flan-t5-small")
@app.middleware("http")
async def log_requests(request: Request, call_next):
print(f"Incoming request: {request.method} {request.url}")
response = await call_next(request)
return response
@app.get("/")
def read_root():
return {
"message": "Welcome to the FastAPI text generation API! Use the /generate endpoint with a 'text' query parameter."
}
@app.get("/generate")
def generate(text: str):
try:
output = pipe(text)
return {"output": output[0]["generated_text"]}
except Exception as e:
return JSONResponse(
status_code=500,
content={"error": f"An error occurred: {str(e)}"}
)