deve

app.py

@@ -31,27 +31,32 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
     try:
         torch.cuda.empty_cache()
         torch.nn.CrossEntropyLoss()
-        metric = evaluate.load("rouge", cache_dir='/data/cache')
+        rouge_metric = evaluate.load("rouge", cache_dir='/data/cache')
         def compute_metrics(eval_preds):
             preds, labels = eval_preds
             if isinstance(preds, tuple):
                 preds = preds[0]
+        
             # Replace -100s used for padding as we can't decode them
-           # preds = np.where(preds != -100, preds, tokenizer.pad_token_id)
-            #preds = np.array(preds)
+            preds = np.where(preds != -100, preds, tokenizer.pad_token_id)
+            labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+        
+            # Decode predictions and labels
             decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
-            if isinstance(labels, tuple):
-                labels = labels[0]            
-            #labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
-            #labels = np.array(labels)
             decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
-      
-            result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
+        
+            # Compute ROUGE metrics
+            result = rouge_metric.compute(predictions=decoded_preds, references=decoded_labels)
             result = {k: round(v * 100, 4) for k, v in result.items()}
+        
+            # Calculate accuracy
+            accuracy = accuracy_score(decoded_labels, decoded_preds)
+            result["accuracy"] = round(accuracy * 100, 4)
+        
+            # Calculate average generation length
             prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
             result["gen_len"] = np.mean(prediction_lens)
-            accuracy = accuracy_score(decoded_labels, decoded_preds)
-            result["eval_accuracy"] = round(accuracy * 100, 4)            
+        
             return result
         
         login(api_key.strip())
@@ -226,8 +231,8 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
         trainer = Trainer(
             model=model,
             args=training_args,
-            train_dataset=dataset['train'], 
-            eval_dataset=dataset['validation'], 
+            train_dataset=train_set['train'], 
+            eval_dataset=train_set['validation'], 
             compute_metrics=compute_metrics,
             #data_collator=data_collator,
             processing_class=tokenizer,