Spaces:
Sleeping
Sleeping
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. | |
def generate(prompt: str): | |
output = pipe(prompt) | |
return {'output': output[0]['generated_text']} | |