Update app.py

app.py

@@ -5,7 +5,7 @@ import pytesseract
 import numpy as np
 from flask import Flask, request, jsonify
 from flask_cors import CORS
-import transformers  # Full import for logging
+import transformers
 from transformers import PegasusForConditionalGeneration, PegasusTokenizer, BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments
 from datasets import load_dataset, concatenate_datasets
 import torch
@@ -20,7 +20,6 @@ BERT_MODEL_DIR = 'fine_tuned_bert'
 LEGALBERT_MODEL_DIR = 'fine_tuned_legalbert'
 MAX_FILE_SIZE = 100 * 1024 * 1024
 
-# Ensure upload folder exists
 if not os.path.exists(UPLOAD_FOLDER):
     os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 
@@ -39,14 +38,20 @@ def load_or_finetune_pegasus():
         model = PegasusForConditionalGeneration.from_pretrained("google/pegasus-xsum")
         
         cnn_dm = load_dataset("cnn_dailymail", "3.0.0", split="train[:5000]")
-        xsum = load_dataset("xsum", split="train[:5000]", trust_remote_code=True)  # Added trust_remote_code=True
+        xsum = load_dataset("xsum", split="train[:5000]", trust_remote_code=True)
         combined_dataset = concatenate_datasets([cnn_dm, xsum])
         
         def preprocess_function(examples):
-            inputs = tokenizer(examples["article"] if "article" in examples else examples["document"], 
-                              max_length=512, truncation=True, padding="max_length")
-            targets = tokenizer(examples["highlights"] if "highlights" in examples else examples["summary"], 
-                               max_length=400, truncation=True, padding="max_length")
+            # Extract the correct text field (article or document) as a list of strings
+            texts = [examples["article"][i] if "article" in examples else examples["document"][i] 
+                     for i in range(len(examples["article"] if "article" in examples else examples["document"]))]
+            inputs = tokenizer(texts, max_length=512, truncation=True, padding="max_length", return_tensors="pt")
+            
+            # Extract the correct summary field (highlights or summary) as a list of strings
+            summaries = [examples["highlights"][i] if "highlights" in examples else examples["summary"][i] 
+                         for i in range(len(examples["highlights"] if "highlights" in examples else examples["summary"]))]
+            targets = tokenizer(summaries, max_length=400, truncation=True, padding="max_length", return_tensors="pt")
+            
             inputs["labels"] = targets["input_ids"]
             return inputs