Update app.py

app.py

@@ -15,9 +15,9 @@ from sklearn.metrics.pairwise import cosine_similarity
 app = Flask(__name__)
 CORS(app)
 UPLOAD_FOLDER = os.path.join(os.getcwd(), 'uploads')
-PEGASUS_MODEL_DIR = 'fine_tuned_pegasus'
-BERT_MODEL_DIR = 'fine_tuned_bert'
-LEGALBERT_MODEL_DIR = 'fine_tuned_legalbert'
+PEGASUS_MODEL_DIR = '/app/fine_tuned_pegasus'
+BERT_MODEL_DIR = '/app/fine_tuned_bert'
+LEGALBERT_MODEL_DIR = '/app/fine_tuned_legalbert'
 MAX_FILE_SIZE = 100 * 1024 * 1024
 
 if not os.path.exists(UPLOAD_FOLDER):
@@ -37,13 +37,11 @@ def load_or_finetune_pegasus():
         tokenizer = PegasusTokenizer.from_pretrained("google/pegasus-xsum")
         model = PegasusForConditionalGeneration.from_pretrained("google/pegasus-xsum")
         
-        # Load and normalize datasets
         cnn_dm = load_dataset("cnn_dailymail", "3.0.0", split="train[:5000]").rename_column("article", "text").rename_column("highlights", "summary")
         xsum = load_dataset("xsum", split="train[:5000]", trust_remote_code=True).rename_column("document", "text")
         combined_dataset = concatenate_datasets([cnn_dm, xsum])
         
         def preprocess_function(examples):
-            # Directly use normalized 'text' and 'summary' fields
             inputs = tokenizer(examples["text"], max_length=512, truncation=True, padding="max_length", return_tensors="pt")
             targets = tokenizer(examples["summary"], max_length=400, truncation=True, padding="max_length", return_tensors="pt")
             inputs["labels"] = targets["input_ids"]
@@ -54,7 +52,7 @@ def load_or_finetune_pegasus():
         eval_dataset = tokenized_dataset.select(range(8000, 10000))
         
         training_args = TrainingArguments(
-            output_dir="./pegasus_finetune",
+            output_dir="/app/pegasus_finetune",
             num_train_epochs=3,
             per_device_train_batch_size=1,
             per_device_eval_batch_size=1,
@@ -117,7 +115,7 @@ def load_or_finetune_bert():
         eval_dataset = tokenized_dataset.select(range(int(0.8 * len(tokenized_dataset)), len(tokenized_dataset)))
         
         training_args = TrainingArguments(
-            output_dir="./bert_finetune",
+            output_dir="/app/bert_finetune",
             num_train_epochs=3,
             per_device_train_batch_size=8,
             per_device_eval_batch_size=8,
@@ -180,7 +178,7 @@ def load_or_finetune_legalbert():
         eval_dataset = tokenized_dataset.select(range(int(0.8 * len(tokenized_dataset)), len(tokenized_dataset)))
         
         training_args = TrainingArguments(
-            output_dir="./legalbert_finetune",
+            output_dir="/app/legalbert_finetune",
             num_train_epochs=3,
             per_device_train_batch_size=8,
             per_device_eval_batch_size=8,