Update train_engine.py

train_engine.py

@@ -5,62 +5,80 @@ from transformers import (
     Trainer, 
     TrainingArguments
 )
-from datasets import load_dataset
+from datasets import load_dataset, Dataset
 import pandas as pd
 import os
 
 def train_on_devign(base_model="microsoft/codebert-base", output_dir="./trained_model"):
-    print(f"🚀 Initializing Autotrain Engine for {base_model}")
+    print(f"🚀 Initializing Autotrain Engine (Precision v2) for {base_model}")
     
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"💻 Using hardware: {device}")
+
     # 1. Load specialized Devign dataset
     print("📥 Loading Devign dataset from Hugging Face Hub...")
     try:
-        dataset = load_dataset("DetectVul/devign")
+        remote_data = load_dataset("DetectVul/devign", split="train[:5000]") # Limit to 5k for speed
     except Exception as e:
         print(f"Failed to load Devign: {e}. Falling back to sample dataset.")
         return
 
+    # 2. Integrate Local Feedback Data (Active Learning)
+    feedback_file = "feedback_dataset.csv"
+    if os.path.exists(feedback_file):
+        print("📈 Merging local feedback data into training set...")
+        fb_df = pd.read_csv(feedback_file)
+        # Assuming CSV has 'original_code' and we treat applied fixes as 'Safe' (Label 0) or similar
+        # For simplicity, we just add the code and label it
+        fb_data = Dataset.from_pandas(fb_df.rename(columns={'original_code': 'func'}))
+        # Add labels if missing
+        if 'label' not in fb_data.column_names:
+            fb_data = fb_data.add_column("label", [1] * len(fb_data)) # Treat feedback items as vulnerable patterns we should recognize
+        
+        # Merge remote and local
+        from datasets import concatenate_datasets
+        dataset = concatenate_datasets([remote_data, fb_data])
+    else:
+        dataset = remote_data
+
     tokenizer = AutoTokenizer.from_pretrained(base_model)
 
     def tokenize_function(examples):
         return tokenizer(examples["func"], padding="max_length", truncation=True, max_length=512)
 
-    print("✂️ Tokenizing dataset...")
+    print("✂️ Tokenizing hybrid dataset...")
     tokenized_datasets = dataset.map(tokenize_function, batched=True)
 
-    # 2. Load Model
+    # 3. Load Model
     print("🧠 Loading Base Model...")
-    model = AutoModelForSequenceClassification.from_pretrained(base_model, num_labels=2)
+    model = AutoModelForSequenceClassification.from_pretrained(base_model, num_labels=2).to(device)
 
-    # 3. Setup Training
+    # 4. Setup Training
     training_args = TrainingArguments(
         output_dir=output_dir,
-        evaluation_strategy="epoch",
-        learning_rate=2e-5,
-        per_device_train_batch_size=8, # Optimized for high-performance
-        per_device_eval_batch_size=8,
         num_train_epochs=3,
+        per_device_train_batch_size=4, # Reduced for stability on wider range of hardware
+        learning_rate=2e-5,
         weight_decay=0.01,
-        push_to_hub=False,
         logging_dir='./logs',
+        save_strategy="no",
+        report_to="none"
     )
 
     trainer = Trainer(
         model=model,
         args=training_args,
-        train_dataset=tokenized_datasets["train"],
-        eval_dataset=tokenized_datasets["test"],
+        train_dataset=tokenized_datasets,
     )
 
-    # 4. Train
-    print("🔥 Starting Fine-tuning cycle...")
+    # 5. Train
+    print("🔥 Starting active learning cycle...")
     trainer.train()
 
-    # 5. Save & Update
-    print(f"✅ Training Complete. Saving to {output_dir}")
+    # 6. Save results
+    print(f"✅ Training Complete. Saving weights to {output_dir}")
     model.save_pretrained(output_dir)
     tokenizer.save_pretrained(output_dir)
 
 if __name__ == "__main__":
-    # In a real scenario, this would be triggered by /train
     train_on_devign()