Create train.py

train.py

@@ -0,0 +1,68 @@
+import torch
+import json
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
+from peft import LoraConfig, get_peft_model
+from datasets import Dataset
+
+# ✅ Load Extracted Data
+with open("medical_dataset.json", "r", encoding="utf-8") as f:
+    data = json.load(f)
+
+dataset = Dataset.from_list(data)
+
+# ✅ Load Tokenizer
+model_name = "meta-llama/Llama-2-7b-hf"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# ✅ Tokenize Data
+def preprocess_function(examples):
+    inputs = [f"Medical Q&A: {ex['prompt']} {ex['response']}" for ex in examples]
+    model_inputs = tokenizer(inputs, padding="max_length", truncation=True, max_length=512)
+    model_inputs["labels"] = model_inputs["input_ids"].copy()
+    return model_inputs
+
+tokenized_dataset = dataset.map(preprocess_function, batched=True)
+tokenized_dataset = tokenized_dataset.remove_columns(["prompt", "response"])
+
+# ✅ Load Model with QLoRA (4-bit Precision)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    load_in_4bit=True, 
+    device_map="auto"
+)
+
+lora_config = LoraConfig(
+    r=16,
+    lora_alpha=32,
+    target_modules=["q_proj", "v_proj"],
+    lora_dropout=0.05,
+    bias="none"
+)
+
+model = get_peft_model(model, lora_config)
+
+# ✅ Define Training Arguments
+training_args = TrainingArguments(
+    output_dir="./medical_llama2",
+    per_device_train_batch_size=1,
+    num_train_epochs=2,  # 2 Epochs
+    logging_dir="./logs",
+    save_steps=100,
+    evaluation_strategy="no"
+)
+
+# ✅ Train Model
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_dataset
+)
+
+trainer.train()
+
+# ✅ Save Model
+trainer.save_model("fine_tuned_medical_llama2")
+tokenizer.save_pretrained("fine_tuned_medical_llama2")