File size: 924 Bytes
efb7634 |
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 30 31 32 33 34 35 36 37 |
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"]}
|