gnharishkumar's picture
Upload 3 files
efb7634 verified
raw
history blame
924 Bytes
from fastapi import FastAPI
from transformers import pipeline
# create FastAPI app
app = FastAPI()
# create a text-generation pipeline
pipe = pipeline("text2text-generation", model="google/flan-t5-small")
@app.get("/")
def home():
"""
Home route for the FastAPI app.
Returns:
dict: A dictionary with a message indicating that it is a simple FastAPI app for text generation using T5.
"""
return {"message": "This is a simple FastAPI app for text generation using T5"}
# create a route for text generation
@app.get("/generate/")
def generate_text(prompt: str):
"""
Route for generating text using the T5 model.
Args:
prompt (str): The prompt for the text generation.
Returns:
dict: A dictionary with the generated text.
"""
output = pipe(text=prompt)
return {"generated_text": pipe(prompt)[0]["generated_text"]}