Update app.py

app.py

@@ -1,11 +1,7 @@
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
+import gradio as gr
 from transformers import BartTokenizer, BartForConditionalGeneration
 import torch
 
-# Initialize the FastAPI app
-app = FastAPI()
-
 # Load the fine-tuned BART model and tokenizer from the local directory
 MODEL_DIR = './BART model small/model'
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -13,17 +9,12 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 tokenizer = BartTokenizer.from_pretrained(MODEL_DIR)
 model = BartForConditionalGeneration.from_pretrained(MODEL_DIR).to(device)
 
-# Define a request model for the API input
-class Article(BaseModel):
-    text: str
-
-# API Endpoint for summarization
-@app.post("/summarize")
-async def summarize(article: Article):
+# Define the summarization function
+def summarize(text):
     try:
         # Tokenize the input article
         inputs = tokenizer(
-            article.text,
+            text,
             return_tensors="pt",
             max_length=1024,
             truncation=True
@@ -42,8 +33,20 @@ async def summarize(article: Article):
         # Decode the summary
         summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
 
-        # Return the summary in the response
-        return {"summary": summary}
+        return summary
     
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        return str(e)
+
+# Create Gradio interface
+# Textbox input for the article and output for the summary
+interface = gr.Interface(
+    fn=summarize,          # The function to summarize the article
+    inputs="text",         # Input is a text box where users can input the article text
+    outputs="text",        # Output is a text box displaying the summary
+    title="BART Summarization",  # The title of the app
+    description="Enter an article to generate a summary using a fine-tuned BART model."
+)
+
+# Launch the Gradio app
+interface.launch()