Third_project / app.py
arshadrana's picture
Update app.py
bef1f90 verified
raw
history blame contribute delete
565 Bytes
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']}