File size: 697 Bytes
f93b405
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from fastapi import FastAPI
from transformers import pipeline

# Create a FastAPI instance
app = FastAPI()

# The root endpoint of the application
@app.get("/")
def root():
    return {'message': 'Welcome to the Text Generation API!'}

# Initialize the text generation pipeline
pipe = pipeline('text2test-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']}