Update app.py

app.py

@@ -1,28 +1,44 @@
 import gradio as gr
 from transformers import LEDForConditionalGeneration, LEDTokenizer
 import torch
+from datasets import load_dataset
 
-# Load the model and tokenizer
-model = LEDForConditionalGeneration.from_pretrained("./summary_generation_Led_4")
+# Set device to GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load the LED model and tokenizer
+model = LEDForConditionalGeneration.from_pretrained("./summary_generation_Led_4").to(device)
 tokenizer = LEDTokenizer.from_pretrained("./summary_generation_Led_4")
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model = model.to(device)
+# Normalize the input text (plot synopsis)
+def normalize_text(text):
+    text = text.lower()  # Lowercase the text
+    text = re.sub(r'\s+', ' ', text).strip()  # Remove extra spaces and newlines
+    text = re.sub(r'[^\w\s]', '', text)  # Remove non-alphanumeric characters
+    return text
 
-# Define the function for generating summaries
+# Function to preprocess and generate summaries
 def generate_summary(plot_synopsis):
-    inputs = tokenizer(plot_synopsis, max_length=3000, truncation=True, padding="max_length", return_tensors="pt")
+    # Preprocess the plot_synopsis
+    inputs = tokenizer("summarize: " + normalize_text(plot_synopsis), 
+                       max_length=3000, truncation=True, padding="max_length", return_tensors="pt")
     inputs = inputs.to(device)
-    outputs = model.generate(inputs['input_ids'], max_length=315, min_length=20, length_penalty=2.0, num_beams=4, early_stopping=True)
+
+    # Generate the summary
+    outputs = model.generate(inputs["input_ids"], max_length=315, min_length=20, 
+                             length_penalty=2.0, num_beams=4, early_stopping=True)
+    
     summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return summary
 
-# Create a Gradio interface
-interface = gr.Interface(fn=generate_summary,
-                         inputs=gr.Textbox(label="Plot Synopsis", lines=10, placeholder="Enter plot synopsis here..."),
-                         outputs=gr.Textbox(label="Plot Summary"),
-                         title="Plot Summary Generator",
-                         description="This demo generates a plot summary based on a plot synopsis using a fine-tuned LED model.")
+# Gradio interface to take plot synopsis and output a generated summary
+interface = gr.Interface(
+    fn=generate_summary,
+    inputs=gr.Textbox(label="Plot Synopsis", lines=10, placeholder="Enter the plot synopsis here..."),
+    outputs=gr.Textbox(label="Generated Summary"),
+    title="Plot Summary Generator",
+    description="This demo generates a plot summary based on the plot synopsis using a fine-tuned LED model."
+)
 
-# Launch the interface
+# Launch the Gradio interface
 interface.launch()