text-generation / app.py
lucapantea
swagger docs
140ff41
raw
history blame contribute delete
581 Bytes
from fastapi import FastAPI
from transformers import pipeline
# Create a FastAPI instance
app = FastAPI(docs_url='/')
# Initialize the text generation pipeline
pipe = pipeline('text2text-generation', model='google/flan-t5-small')
# The GET endpoint of the application,
# which corresponds to the text generation
# functionality. The generate() method takes
# a prompt as input and returns the generated
# text as output in the form of a JSON object.
@app.get("/generate")
def generate(prompt: str):
output = pipe(prompt)
return {'output': output[0]['generated_text']}