dev

app.py

@@ -4,6 +4,9 @@ from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelFor
 from transformers import DataCollatorForSeq2Seq
 from datasets import load_dataset, concatenate_datasets, load_from_disk
 import traceback
+from sklearn.metrics import accuracy_score
+import numpy as np
+
 import os
 from huggingface_hub import login
 from peft import get_peft_model, LoraConfig
@@ -12,7 +15,15 @@ os.environ['HF_HOME'] = '/data/.huggingface'
 
 @spaces.GPU(duration=120)
 def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
-    try:        
+    try:
+        def compute_metrics(eval_pred):
+            logits, labels = eval_pred
+            predictions = np.argmax(logits, axis=1)
+            accuracy = accuracy_score(labels, predictions)
+            return {
+                'eval_accuracy': accuracy,
+                'eval_loss': eval_pred.loss,  # If you want to include loss as well
+            }        
         login(api_key.strip())
         lora_config = LoraConfig(
             r=16,  # Rank of the low-rank adaptation
@@ -23,7 +34,7 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
     
         # Load the model and tokenizer
         model = AutoModelForSeq2SeqLM.from_pretrained(model_name.strip(), num_labels=2)
-        model = get_peft_model(model, lora_config)
+        #model = get_peft_model(model, lora_config)
         
     
         # Set training arguments
@@ -58,6 +69,7 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
         max_length = 128
         try:
             tokenized_train_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset')
+            print(tokenized_train_dataset[0])
             tokenized_test_dataset = load_from_disk(f'/data/{hub_id.strip()}_test_dataset')
             
             # Create Trainer
@@ -66,6 +78,7 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
                 args=training_args,
                 train_dataset=tokenized_train_dataset,
                 eval_dataset=tokenized_test_dataset,
+                compute_metrics=compute_metrics,
                 #callbacks=[LoggingCallback()], 
             )            
         except:
@@ -107,6 +120,7 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
                 args=training_args,
                 train_dataset=tokenized_datasets['train'],
                 eval_dataset=tokenized_datasets['test'],
+                compute_metrics=compute_metrics,
                 #callbacks=[LoggingCallback()], 
             )            
     
@@ -140,7 +154,7 @@ try:
             gr.Textbox(label="HF hub to push to after training"),
             gr.Textbox(label="HF API token"),
             gr.Slider(minimum=1, maximum=10, value=3, label="Number of Epochs", step=1),
-            gr.Slider(minimum=1, maximum=500, value=1, label="Batch Size", step=1),
+            gr.Slider(minimum=1, maximum=2000, value=1, label="Batch Size", step=1),
             gr.Slider(minimum=1, maximum=1000, value=1, label="Learning Rate (e-5)", step=1),
             gr.Slider(minimum=1, maximum=100, value=1, label="Gradient accumulation", step=1), 
         ],

requirements.txt

@@ -3,3 +3,5 @@ transformers
 datasets
 peft
 huggingface_hub
+scikit-learn
+numpy