Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import shutil
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import pandas as pd
 import torch
 import matplotlib.pyplot as plt
@@ -9,7 +9,7 @@ import seaborn as sns
 import base64
 
 # Define constants
-MODEL_NAME = "gpt2"  # Publicly accessible model suitable for CPU
+MODEL_NAME = "facebook/bart-large-cnn"  # Fine-tuned for summarization
 FIGURES_DIR = "./figures"
 EXAMPLE_DIR = "./example"
 EXAMPLE_FILE = os.path.join(EXAMPLE_DIR, "titanic.csv")
@@ -36,7 +36,7 @@ if not os.path.isfile(EXAMPLE_FILE):
 print("Loading model and tokenizer...")
 try:
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
     model.to('cpu')  # Ensure the model runs on CPU
     print("Model and tokenizer loaded successfully.")
 except Exception as e:
@@ -86,18 +86,15 @@ def generate_summary(prompt):
 
     # Generate response
     with torch.no_grad():
-        outputs = model.generate(
+        summary_ids = model.generate(
             inputs,
             max_length=500,
-            do_sample=True,
-            top_p=0.95,
-            temperature=0.7,
-            eos_token_id=tokenizer.eos_token_id,
-            pad_token_id=tokenizer.eos_token_id
+            num_beams=4,
+            early_stopping=True
         )
 
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    return summary
 
 def analyze_data(data_file_path):
     """Perform data analysis on the uploaded CSV file."""
@@ -249,3 +246,5 @@ if __name__ == "__main__":
 
 
 
+
+