Add support for jupyter notebooks for training

tools/train/setup.py

@@ -15,6 +15,8 @@ setup(
         'datasets>=2.14.6',
         'einops>=0.7.0',
         'flake8>=4.0.0',
+        'ipykernel>=6.0.0', 
+        'ipywidgets>=8.0.0',
         'peft>=0.8.0',
         'peft>=0.8.0',
         'pytest-cov>=4.0.0',

tools/train/train/qwq.py

@@ -11,13 +11,13 @@ from dataclasses import dataclass
 from pathlib import Path
 from transformers import TrainingArguments
 from trl import SFTTrainer, DataCollatorForCompletionOnlyLM
+from typing import Optional, List
 import argparse
 import os
-import torch
 
 @dataclass
 class Args:
-    def __init__(self):
+    def __init__(self, args: Optional[List[str]]):
         parser = argparse.ArgumentParser(description='Train SIA model using QwQ')
         parser.add_argument(
             '--config',
@@ -43,7 +43,10 @@ class Args:
             default=os.environ.get('SIA_HF_API_KEY'),
             help='HuggingFace API key'
         )
-        self.args = parser.parse_args()
+        if args is None:
+            self.args = parser.parse_args()
+        else:
+            self.args = parser.parse_args(args)
 
     @property
     def config_path(self) -> Path:
@@ -86,8 +89,12 @@ def main():
             "gate_proj", "up_proj", "down_proj",
         ], # Remove QKVO if out of memory
         lora_alpha = lora_rank,
-        use_gradient_checkpointing = "unsloth", # Enable long context finetuning
+        lora_dropout = 0, # Supports any, but = 0 is optimized
+        bias = "none",    # Supports any, but = "none" is optimized
+        use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
         random_state = 3407,
+        use_rslora = False,  # We support rank stabilized LoRA
+        loftq_config = None, # And LoftQ
     )
 
     response_template = tokenizer.apply_chat_template(