Update train.py

train.py

@@ -0,0 +1,55 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
+from datasets import load_dataset
+from peft import LoraConfig, get_peft_model
+
+# Model & Tokenizer
+MODEL_NAME = "TinyLlama/TinyLlama-1.1B"  # Adjust if using your own model
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16)
+
+# Apply LoRA for Efficient Fine-Tuning
+peft_config = LoraConfig(
+    r=8,  # Low-rank adaptation size
+    lora_alpha=16,
+    lora_dropout=0.05,
+    bias="none",
+    task_type="CAUSAL_LM"
+)
+model = get_peft_model(model, peft_config)
+
+# Load Dataset (OASST1)
+dataset = load_dataset("OpenAssistant/oasst1", split="train[:10%]")  # Use 10% of dataset
+
+# Tokenization Function
+def tokenize_function(examples):
+    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)
+
+# Tokenize Dataset
+tokenized_datasets = dataset.map(tokenize_function, batched=True)
+
+# Training Arguments
+training_args = TrainingArguments(
+    output_dir="./tinyllama-finetuned",
+    evaluation_strategy="epoch",
+    save_strategy="epoch",
+    per_device_train_batch_size=2,  # Adjust for CPU
+    per_device_eval_batch_size=2,
+    num_train_epochs=3,
+    logging_dir="./logs",
+    report_to="none"
+)
+
+# Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_datasets,
+)
+
+# Start Training
+trainer.train()
+
+# Save Fine-Tuned Model
+model.save_pretrained("./tinyllama-finetuned")
+tokenizer.save_pretrained("./tinyllama-finetuned")