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']}