Update fine_tuning.py

fine_tuning.py

@@ -10,7 +10,6 @@ import transformers
 from transformers import AutoModelForCausalLM, DataCollatorForLanguageModeling, Trainer, TrainingArguments
 from peft import LoraConfig, get_peft_model
 from sentence_transformers import SentenceTransformer, util
-import spaces
 
 # -----------------------------
 # ENVIRONMENT / CACHE
@@ -51,7 +50,6 @@ class LoraLinear(nn.Module):
         else:
             self.lora_A, self.lora_B, self.lora_dropout = None, None, None
 
-    @spaces.GPU
     def forward(self, x):
         result = F.linear(x, self.weight, self.bias)
         if self.r > 0:
@@ -71,7 +69,6 @@ class MoELoRALinear(nn.Module):
         ])
         self.gate = nn.Linear(base_linear.in_features, num_experts)
 
-    @spaces.GPU
     def forward(self, x):
         base_out = self.base_linear(x)
         gate_scores = torch.softmax(self.gate(x), dim=-1)
@@ -80,7 +77,6 @@ class MoELoRALinear(nn.Module):
             expert_out += gate_scores[..., i:i+1] * expert(x)
         return base_out + expert_out
 
-@spaces.GPU
 def replace_proj_with_moe_lora(model, r=8, num_experts=2, k=1, lora_alpha=16, lora_dropout=0.05):
     for layer in model.model.layers:
         for proj_name in ["up_proj", "down_proj"]:
@@ -113,7 +109,6 @@ def preprocess(example):
 # -----------------------------
 # LOAD & TRAIN MODEL
 # -----------------------------
-@spaces.GPU
 def load_and_train(model_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0"):
     global model 
     current_dir = os.path.dirname(os.path.abspath(__file__))
@@ -144,7 +139,7 @@ def load_and_train(model_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0"):
     data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
     training_args = TrainingArguments(
-        learning_rate=5e-5,
+        learning_rate=1e-4,
         lr_scheduler_type="constant",
         output_dir="./results",
         num_train_epochs=4,
@@ -206,7 +201,6 @@ def validate_query(query: str, threshold: float = 0.5) -> bool:
 # -----------------------------
 # GENERATE ANSWER
 # -----------------------------
-@spaces.GPU
 def generate_answer(prompt, max_tokens=200):
     if prompt.strip() == "":
         return "Please enter a prompt!"