File size: 565 Bytes
a993607
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bef1f90
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
from transformers import pipeline
from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI()

# Load the text generation pipeline from Hugging Face
generator = pipeline('text-generation', model='gpt2')

# Define the request body
class TextInput(BaseModel):
    input_text: str

@app.post("/generate")
def generate_text(text_input: TextInput):
    # Generate text from the model based on the input
    generated = generator(text_input.input_text, max_length=150, num_return_sequences=1)
    return {"generated_text": generated[0]['generated_text']}