dev

app.py

@@ -1,21 +1,190 @@
+import spaces
 import gradio as gr
-from transformers import AutoModelForSeq2SeqLM
+from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM
 from transformers import DataCollatorForSeq2Seq, AutoConfig
+from datasets import load_dataset, concatenate_datasets, load_from_disk
+import traceback
+from sklearn.metrics import accuracy_score
+import numpy as np
+import torch
+import os
+import evaluate
+from huggingface_hub import login
+from peft import get_peft_model, LoraConfig
 
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-print(f"Successfully loaded the model without gradio or spaces, model object: {model}")
+os.environ['HF_HOME'] = '/data/.huggingface'
+'''
+lora_config = LoraConfig(
+    r=16,  # Rank of the low-rank adaptation
+    lora_alpha=32,  # Scaling factor
+    lora_dropout=0.1,  # Dropout for LoRA layers
+    bias="none"  # Bias handling
+)
+model = AutoModelForSeq2SeqLM.from_pretrained('google/t5-efficient-tiny', num_labels=2, force_download=True)
+model = get_peft_model(model, lora_config)
+model.gradient_checkpointing_enable()
+model_save_path = '/data/lora_finetuned_model'  # Specify your desired save path
+model.save_pretrained(model_save_path)
+'''
 
-@spaces.GPU(duration=120)
-def run_train(model_name, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
+def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
+    try:
+        torch.nn.CrossEntropyLoss()
+        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)
+            decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
+            labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+            decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+      
+            result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
+            result = {k: round(v * 100, 4) for k, v in result.items()}
+            prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
+            result["gen_len"] = np.mean(prediction_lens)
+            return result
+        
+        login(api_key.strip())
+   
+    
+        # Load the model and tokenizer
+             
+        
+    
+        # Set training arguments
+        training_args = TrainingArguments(
+            output_dir='/data/results',
+            eval_strategy="steps",  # Change this to steps
+            save_strategy='steps',
+            learning_rate=lr*0.00001,
+            per_device_train_batch_size=int(batch_size),
+            per_device_eval_batch_size=int(batch_size), 
+            num_train_epochs=int(num_epochs),
+            weight_decay=0.01,
+            #gradient_accumulation_steps=int(grad),
+            #max_grad_norm = 1.0, 
+            load_best_model_at_end=True,
+            metric_for_best_model="accuracy",
+            greater_is_better=True,
+            logging_dir='/data/logs',
+            logging_steps=10,
+            #push_to_hub=True,
+            hub_model_id=hub_id.strip(),
+            fp16=True,
+            #lr_scheduler_type='cosine',
+            save_steps=100,  # Save checkpoint every 500 steps
+            save_total_limit=3, 
+        )
+        # Check if a checkpoint exists and load it
+        if os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir):
+            print("Loading model from checkpoint...")
+            model = AutoModelForSeq2SeqLM.from_pretrained(training_args.output_dir)        
+    
+        def tokenize_function(examples):
+            
+            # Assuming 'text' is the input and 'target' is the expected output
+            model_inputs = tokenizer(
+                examples['text'], 
+                max_length=max_length,  # Set to None for dynamic padding
+                truncation=True,
+                padding='longest',
+                return_tensors='pt', 
+            )
+        
+            # Setup the decoder input IDs (shifted right)
+            labels = tokenizer(
+                examples['target'], 
+                max_length=max_length,  # Set to None for dynamic padding
+                truncation=True,
+                padding='longest',
+                #text_target=examples['target'],
+                return_tensors='pt', 
+            )
+        
+            # Add labels to the model inputs
+            model_inputs["labels"] = labels["input_ids"]
+            return model_inputs
+    
+        #max_length = 512
+        # Load the dataset
+        dataset = load_dataset(dataset_name.strip())          
+        train_size = len(dataset['train'])
+        half_size = train_size // 2  
+        max_length = model.get_input_embeddings().weight.shape[0]
+        try:
+            tokenized_first_half = load_from_disk(f'/data/{hub_id.strip()}_train_dataset')
+            second_half = dataset['train'].select(range(half_size, train_size))
+            tokenized_second_half = tokenize_function(second_half.to_dict())
+            tokenized_train_dataset = concatenate_datasets([tokenized_first_half, tokenized_second_half])
+            tokenized_test_dataset = tokenize_function(dataset['test'])
+            
+            # Create Trainer
+            trainer = Trainer(
+                model=model,
+                args=training_args,
+                train_dataset=tokenized_train_dataset,
+                eval_dataset=tokenized_test_dataset,
+                compute_metrics=compute_metrics,
+            )            
+        except:                                            
+            tokenizer = AutoTokenizer.from_pretrained('google/t5-efficient-tiny-nh8')
+            # Tokenize the dataset                    
+            first_half = dataset['train'].select(range(half_size))
+            tokenized_half = tokenize_function(first_half.to_dict())
+            
+            tokenized_half.save_to_disk(f'/data/{hub_id.strip()}_train_dataset')                           
+        
+            return 'RUN AGAIN TO LOAD REST OF DATA'
+
+        # Fine-tune the model
+        if os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir):
+            train_result = trainer.train(resume_from_checkpoint=True)
+        else:
+            train_result = trainer.train()
+        trainer.push_to_hub(commit_message="Training complete!")
+    except Exception as e:
+        return f"An error occurred: {str(e)}, TB: {traceback.format_exc()}"
+    return 'DONE!'#train_result
+'''
+# Define Gradio interface
+def predict(text):
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name.strip(), num_labels=2)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    outputs = model(inputs)
+    predictions = outputs.logits.argmax(dim=-1)
+    return predictions.item()
+'''
 
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-    return "WORKS"
+@spaces.GPU(duration=120)
+def run_train(dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
+    def initialize_weights(model):
+        for name, param in model.named_parameters():
+            if 'encoder.block.0.layer.0.DenseReluDense.wi.weight' in name:  # Example layer
+                torch.nn.init.xavier_uniform_(param.data)  # Xavier initialization
+            elif 'encoder.block.0.layer.0.DenseReluDense.wo.weight' in name:  # Another example layer
+                torch.nn.init.kaiming_normal_(param.data)  # Kaiming initialization
+    
+    config = AutoConfig.from_pretrained("google/t5-efficient-tiny")
+    model = AutoModelForSeq2SeqLM.from_config(config)
+    initialize_weights(model)
+    lora_config = LoraConfig(
+        r=16,  # Rank of the low-rank adaptation
+        lora_alpha=32,  # Scaling factor
+        lora_dropout=0.1,  # Dropout for LoRA layers
+        bias="none"  # Bias handling
+    )
+    model = get_peft_model(model, lora_config)
+    result = fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad)
+    return result
 # Create Gradio interface
 try:    
     iface = gr.Interface(
         fn=run_train,
         inputs=[
-            gr.Textbox(label="Model Name (e.g., 'google/t5-efficient-tiny-nh8')"),
             gr.Textbox(label="Dataset Name (e.g., 'imdb')"),
             gr.Textbox(label="HF hub to push to after training"),
             gr.Textbox(label="HF API token"),
@@ -28,8 +197,19 @@ try:
         title="Fine-Tune Hugging Face Model",
         description="This interface allows you to fine-tune a Hugging Face model on a specified dataset."
     )
-
+    '''
+    iface = gr.Interface(
+        fn=predict,
+        inputs=[
+            gr.Textbox(label="Query"),
+        ],
+        outputs="text",
+        title="Fine-Tune Hugging Face Model",
+        description="This interface allows you to test a fine-tune Hugging Face model."
+    )
+    '''
     # Launch the interface
     iface.launch()    
 except Exception as e:
-    print(f"An error occurred: {str(e)}, TB: {traceback.format_exc()}")
+    print(f"An error occurred: {str(e)}, TB: {traceback.format_exc()}")
+