added finetuning and trainer scripts

finetuning.py

@@ -0,0 +1,70 @@
+from datasets import load_dataset
+from trl import SFTTrainer
+from peft import LoraConfig
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments,  BitsAndBytesConfig
+
+#Check if you do not have any import issue to use the Fast Mamba Kernel
+#Will (very appropriately) break before loading the weights.
+import mamba_ssm
+
+#With 4bit quants have to manually correct modeling_jamba.py on l. 1070:
+#if not is_fast_path_available or "cuda" not in self.x_proj.weight.device.type:
+#becoming:
+#if not is_fast_path_available:
+
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    llm_int4_skip_modules=["mamba"] #Maybe not necessary (per axoltl) but to test.
+)
+
+tokenizer = AutoTokenizer.from_pretrained("jamba")
+
+dataset = load_dataset("VishnuPJ/Malayalam_CultureX_IndicCorp_SMC", split="train")
+training_args = TrainingArguments(
+    output_dir="./results",
+    num_train_epochs=1,
+    per_device_train_batch_size=1,
+    gradient_accumulation_steps=4,
+    optim = "adamw_8bit",
+    max_grad_norm = 0.3,
+    weight_decay = 0.001,
+    warmup_ratio = 0.03,
+    gradient_checkpointing=True,
+    logging_dir='./logs',
+    logging_steps=1,
+    max_steps=50,
+    group_by_length=True,
+    lr_scheduler_type = "linear",
+    learning_rate=2e-3
+)
+lora_config = LoraConfig(
+    lora_alpha=16,
+    lora_dropout=0.05,
+    init_lora_weights=False,
+    r=8,
+    target_modules=["embed_tokens", "x_proj", "in_proj", "out_proj"],
+    task_type="CAUSAL_LM",
+    bias="none"
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    "jamba",
+    trust_remote_code=True, 
+    device_map='auto',
+    attn_implementation="flash_attention_2", 
+    quantization_config=quantization_config, 
+    use_mamba_kernels=True
+    )
+
+trainer = SFTTrainer(
+    model=model,
+    tokenizer=tokenizer,
+    args=training_args,
+    peft_config=lora_config,
+    train_dataset=dataset,
+    max_seq_length = 256,
+    dataset_text_field="quote",
+)
+
+trainer.train()

trainer.py

@@ -0,0 +1,38 @@
+# Load model directly
+from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
+
+tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-tiny-random")
+model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-tiny-random",  trust_remote_code=True)
+
+from datasets import load_dataset
+from trl import SFTTrainer
+from peft import LoraConfig
+
+dataset=load_dataset("rajeshradhakrishnan/malayalam_wiki")
+
+training_args = TrainingArguments(
+    output_dir="./results",
+    num_train_epochs=3,
+    per_device_train_batch_size=3,
+    logging_dir='./logs',
+    logging_steps=10,
+    learning_rate=2e-3
+)
+
+lora_config = LoraConfig(
+    r=8,
+    target_modules=["embed_tokens", "x_proj", "in_proj", "out_proj"],
+    task_type="CAUSAL_LM",
+    bias="none"
+)
+
+trainer = SFTTrainer(
+    model=model,
+    tokenizer=tokenizer,
+    args=training_args,
+    peft_config=lora_config,
+    train_dataset=dataset["train"],
+    dataset_text_field="text",
+)
+
+trainer.train()