Upload fine_tune.py

fine_tune.py

@@ -0,0 +1,91 @@
+import torch
+torch.cuda.empty_cache()
+
+# ✅ Clear GPU memory before training
+import torch
+torch.cuda.empty_cache()
+
+# ✅ Load necessary libraries
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, TrainingArguments, Trainer
+import os
+
+# ✅ Load Dataset
+dataset = load_dataset("zainabfatima097/My_Dataset")  # Change to your dataset path
+
+# ✅ Check available splits
+print(f"Available dataset splits: {dataset.keys()}")
+
+# ✅ If dataset has only 'validation' split, rename it to 'train'
+if "train" not in dataset:
+    dataset["train"] = dataset["validation"]
+
+# ✅ Extract Text for Translation Task
+source_lang = "en"
+target_lang = "hi"
+
+def preprocess_function(examples):
+    """ Extracts input and target texts for translation """
+    inputs = [ex[source_lang] for ex in examples["translation"]]
+    targets = [ex[target_lang] for ex in examples["translation"]]
+    return {"input_text": inputs, "target_text": targets}
+
+# ✅ Apply Text Extraction
+dataset = dataset.map(preprocess_function, batched=True)
+
+# ✅ Load Tokenizer
+model_checkpoint = "Helsinki-NLP/opus-mt-en-hi"  # Use your model
+tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
+
+# ✅ Tokenization
+def tokenize_function(examples):
+    inputs = tokenizer(examples["input_text"], truncation=True, padding="max_length", max_length=128)
+    targets = tokenizer(examples["target_text"], truncation=True, padding="max_length", max_length=128)
+    inputs["labels"] = targets["input_ids"]
+    return inputs
+
+# ✅ Apply Tokenization
+tokenized_datasets = dataset.map(tokenize_function, batched=True, remove_columns=["translation", "input_text", "target_text"])
+
+# ✅ Set Train & Validation Splits
+train_dataset = tokenized_datasets["train"]
+eval_dataset = tokenized_datasets.get("validation", train_dataset)  # Use train if validation is missing
+
+# ✅ Load Model
+model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)
+
+# ✅ Training Arguments (Handles Memory Issues)
+training_args = TrainingArguments(
+    output_dir="./results",
+    per_device_train_batch_size=2,   # Reduce batch size to prevent OOM
+    per_device_eval_batch_size=2,
+    gradient_accumulation_steps=4,   # Accumulate gradients to simulate larger batch
+    fp16=True,  # Mixed precision to reduce memory
+    optim="adamw_torch",  # More efficient optimizer
+    evaluation_strategy="epoch",
+    save_strategy="epoch",
+    load_best_model_at_end=True,
+    push_to_hub=False
+)
+
+# ✅ Initialize Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    tokenizer=tokenizer,
+)
+
+# ✅ Train Model (Handling GPU Memory Errors)
+try:
+    trainer.train()
+except torch.cuda.OutOfMemoryError:
+    print("⚠️ CUDA Out of Memory! Switching to CPU...")
+    os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Disable GPU
+    model.to("cpu")
+    trainer.train()
+
+# ✅ Save Model
+trainer.save_model("./final_model")
+print("🎉 Training complete! Model saved.")