Upload final_code.py

final_code.py

@@ -0,0 +1,173 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+from transformers import MT5ForConditionalGeneration, MT5Tokenizer, AdamW
+from transformers import AutoModel, AutoTokenizer
+from sklearn.metrics.pairwise import cosine_similarity
+import pandas as pd
+import matplotlib.pyplot as plt
+import numpy as np
+from huggingface_hub import HfApi, HfFolder, Repository, notebook_login, create_repo, upload_folder
+import os
+import shutil
+
+# ========== CONFIG ==========
+HF_USERNAME = "aarath97"
+HF_REPO = "mt5-dogri-translation"
+MODEL_NAME = "google/mt5-large"
+BATCH_SIZE = 2
+LR = 1e-5
+DPO_STEPS = 100
+HGRL_STEPS = 100
+COMBINED_STEPS = 50
+GAMMA = 3.5
+ALPHA = 0.5
+BETA = 0.5
+
+# ========== LOAD DATA ==========
+df = pd.read_excel("dogri_train.xlsx")
+train_data = list(zip(df['Dogri'], df['English'], df['Unpreffered']))
+
+# ========== TOKENIZERS ==========
+tokenizer = MT5Tokenizer.from_pretrained(MODEL_NAME)
+sbert = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
+sbert_tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
+
+# ========== UTILITIES ==========
+def compute_similarity(sent1, sent2):
+    emb1 = sbert(**sbert_tokenizer(sent1, return_tensors='pt')).last_hidden_state.mean(1)
+    emb2 = sbert(**sbert_tokenizer(sent2, return_tensors='pt')).last_hidden_state.mean(1)
+    return cosine_similarity(emb1.detach().numpy(), emb2.detach().numpy())[0][0]
+
+def hyper_gamma_reward(rho):
+    return rho * np.exp(-GAMMA * (1 - rho))
+
+# ========== DATASET ==========
+class DogriDataset(Dataset):
+    def __init__(self, data):
+        self.data = data
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        return self.data[idx]
+
+dataloader = DataLoader(DogriDataset(train_data), batch_size=BATCH_SIZE, shuffle=True)
+
+# ========== TRAINING ==========
+model = MT5ForConditionalGeneration.from_pretrained(MODEL_NAME).to("cuda")
+optimizer = AdamW(model.parameters(), lr=LR)
+
+dpo_losses, hgrl_losses, final_losses = [], [], []
+
+# ---- DPO Training ----
+for step in range(DPO_STEPS):
+    batch = next(iter(dataloader))
+    loss_batch = []
+    for src, ref, unpref in zip(*batch):
+        input_ids = tokenizer(src, return_tensors='pt', truncation=True, padding=True).input_ids.to("cuda")
+        ref_ids = tokenizer(ref, return_tensors='pt', truncation=True, padding=True).input_ids.to("cuda")
+        unpref_ids = tokenizer(unpref, return_tensors='pt', truncation=True, padding=True).input_ids.to("cuda")
+
+        ref_logprob = model(input_ids=input_ids, labels=ref_ids).loss
+        unpref_logprob = model(input_ids=input_ids, labels=unpref_ids).loss
+
+        logit_diff = -ref_logprob.item() + unpref_logprob.item()
+        beta = 1.0
+        loss = -torch.log(torch.sigmoid(torch.tensor(beta * logit_diff)))
+        loss_batch.append(loss)
+
+    loss_val = torch.stack(loss_batch).mean()
+    loss_val.backward()
+    optimizer.step()
+    optimizer.zero_grad()
+    dpo_losses.append(loss_val.item())
+
+# ---- HGRL Training ----
+for step in range(HGRL_STEPS):
+    batch = next(iter(dataloader))
+    loss_batch = []
+    for src, ref, _ in zip(*batch):
+        input_ids = tokenizer(src, return_tensors='pt').input_ids.to("cuda")
+        gen_ids = model.generate(input_ids)
+        gen_text = tokenizer.decode(gen_ids[0], skip_special_tokens=True)
+
+        rho = compute_similarity(gen_text, ref)
+        reward = hyper_gamma_reward(rho)
+
+        labels = tokenizer(gen_text, return_tensors='pt').input_ids.to("cuda")
+        logprob = model(input_ids=input_ids, labels=labels).loss
+
+        loss = -reward * logprob
+        loss_batch.append(loss)
+
+    loss_val = torch.stack(loss_batch).mean()
+    loss_val.backward()
+    optimizer.step()
+    optimizer.zero_grad()
+    hgrl_losses.append(loss_val.item())
+
+# ---- Combined Training ----
+for step in range(COMBINED_STEPS):
+    batch = next(iter(dataloader))
+    loss_dpo_batch, loss_hgrl_batch = [], []
+    for src, ref, unpref in zip(*batch):
+        input_ids = tokenizer(src, return_tensors='pt').input_ids.to("cuda")
+        ref_ids = tokenizer(ref, return_tensors='pt').input_ids.to("cuda")
+        unpref_ids = tokenizer(unpref, return_tensors='pt').input_ids.to("cuda")
+
+        logprob_ref = model(input_ids=input_ids, labels=ref_ids).loss
+        logprob_unpref = model(input_ids=input_ids, labels=unpref_ids).loss
+        dpo_loss = -torch.log(torch.sigmoid(torch.tensor(logprob_unpref.item() - logprob_ref.item())))
+        loss_dpo_batch.append(dpo_loss)
+
+        gen_ids = model.generate(input_ids)
+        gen_text = tokenizer.decode(gen_ids[0], skip_special_tokens=True)
+        rho = compute_similarity(gen_text, ref)
+        reward = hyper_gamma_reward(rho)
+
+        labels = tokenizer(gen_text, return_tensors='pt').input_ids.to("cuda")
+        logprob = model(input_ids=input_ids, labels=labels).loss
+        hgrl_loss = -reward * logprob
+        loss_hgrl_batch.append(hgrl_loss)
+
+    loss_dpo_mean = torch.stack(loss_dpo_batch).mean()
+    loss_hgrl_mean = torch.stack(loss_hgrl_batch).mean()
+    combined_loss = ALPHA * loss_dpo_mean + BETA * loss_hgrl_mean
+    combined_loss.backward()
+    optimizer.step()
+    optimizer.zero_grad()
+    final_losses.append(combined_loss.item())
+
+# ========== SAVE OUTPUTS ==========
+plt.plot(dpo_losses, label="DPO")
+plt.plot(hgrl_losses, label="HGRL")
+plt.plot(final_losses, label="Combined")
+plt.xlabel("Steps")
+plt.ylabel("Loss")
+plt.legend()
+plt.savefig("loss_curve.png")
+
+with open("loss_report.txt", "w") as f:
+    f.write("DPO Final Loss: {:.4f}\n".format(dpo_losses[-1]))
+    f.write("HGRL Final Loss: {:.4f}\n".format(hgrl_losses[-1]))
+    f.write("Combined Final Loss: {:.4f}\n".format(final_losses[-1]))
+
+# ========== TEST AND SAVE TRANSLATIONS ==========
+test_df = pd.read_excel("in22conv.xlsx")
+test_outputs = []
+for line in test_df.iloc[:, 0].tolist():
+    input_ids = tokenizer(line, return_tensors='pt').input_ids.to("cuda")
+    outputs = model.generate(input_ids)
+    translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    test_outputs.append(translation)
+
+output_df = pd.DataFrame({"Dogri": test_df.iloc[:, 0], "English": test_outputs})
+output_df.to_excel("translated_output.xlsx", index=False)
+
+# ========== PUSH TO HUGGING FACE ==========
+model.save_pretrained("mt5-dogri")
+tokenizer.save_pretrained("mt5-dogri")
+create_repo(f"{HF_USERNAME}/{HF_REPO}", private=False, exist_ok=True)
+upload_folder(repo_id=f"{HF_USERNAME}/{HF_REPO}", folder_path="mt5-dogri")
+print("Model uploaded successfully!")