Upload app.py

app.py

@@ -14,7 +14,17 @@ from peft import get_peft_model, LoraConfig
 
 os.environ['HF_HOME'] = '/data/.huggingface'
 
-@spaces.GPU
+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-nh8', num_labels=2, force_download=True)
+model = get_peft_model(model, lora_config)
+model.gradient_checkpointing_enable()   
+
+@spaces.GPU(duration=120)
 def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
     try:
         torch.cuda.empty_cache()
@@ -27,17 +37,10 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
                 '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
-            lora_alpha=32,  # Scaling factor
-            lora_dropout=0.1,  # Dropout for LoRA layers
-            bias="none"  # Bias handling
-        )        
+   
     
         # Load the model and tokenizer
-        model = AutoModelForSeq2SeqLM.from_pretrained(model_name.strip(), num_labels=2, force_download=True)
-        model = get_peft_model(model, lora_config)
-        model.gradient_checkpointing_enable()        
+             
         
     
         # Set training arguments
@@ -86,7 +89,7 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
         except:
             # Load the dataset
             dataset = load_dataset(dataset_name.strip())
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            tokenizer = AutoTokenizer.from_pretrained('google/t5-efficient-tiny-nh8')
             # Tokenize the dataset
             def tokenize_function(examples):
                 
@@ -148,7 +151,6 @@ try:
     iface = gr.Interface(
         fn=fine_tune_model,
         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"),