Update app.py

app.py

@@ -1,23 +1,18 @@
-import gradio as gr
 from transformers import pipeline
+from fastapi import FastAPI
+from pydantic import BaseModel
 
-# Load a pre-trained text-generation model
-generator = pipeline('text-generation', model='gpt2')
+app = FastAPI()
 
-# Function to process the input text
-def generate_text(input_text):
-    # Generate text using the model
-    result = generator(input_text, max_length=50, num_return_sequences=1)
-    return result[0]['generated_text']
+# Load the text generation pipeline from Hugging Face
+generator = pipeline('text-generation', model='gpt2')
 
-# Set up Gradio interface
-iface = gr.Interface(
-    fn=generate_text,                      # Function to call
-    inputs="text",                         # Input type
-    outputs="text",                        # Output type
-    title="Text Generation API",           # Title of the API
-    description="Provide text as input, and the model will generate some text.",
-)
+# Define the request body
+class TextInput(BaseModel):
+    input_text: str
 
-# Launch the interface
-iface.launch()
+@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=100, num_return_sequences=1)
+    return {"generated_text": generated[0]['generated_text']}