dev

app.py

@@ -6,7 +6,7 @@ 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
 from huggingface_hub import login
 from peft import get_peft_model, LoraConfig
@@ -147,8 +147,17 @@ def predict(text):
 
 @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)
+    print(model.named_parameters())
+    initialize_weights(model)
     lora_config = LoraConfig(
         r=16,  # Rank of the low-rank adaptation
         lora_alpha=32,  # Scaling factor