From b30f952eadd109b0725a371afb5df3cbd5810043 Mon Sep 17 00:00:00 2001
From: 51616 <rujikorn.ch@gmail.com>
Date: Fri, 20 Dec 2024 08:54:03 +0000
Subject: [PATCH] fix acc calc + minor updates

---
 hyperlora/data_utils.py | 14 ++++++++++----
 hyperlora/intx_sft.py   | 27 ++++++++++++++++++---------
 2 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/hyperlora/data_utils.py b/hyperlora/data_utils.py
index 9172622..b5ba802 100644
--- a/hyperlora/data_utils.py
+++ b/hyperlora/data_utils.py
@@ -190,6 +190,7 @@ def get_preprocessing_fn(ds_name):
 #     return tokenize_function
 
 
+# taken from https://github.com/huggingface/trl/issues/632#issuecomment-1972630547
 def get_assistant_start_end_indices(messages, conversation_text):
     """
     Get the start and end indices of the assistant's messages in the conversation text.
@@ -221,14 +222,19 @@ def get_masked_labels(conversation_ids, assistant_ranges):
             yield IGNORE_INDEX
 
 
-def tokenize_chat_messages(example, tokenizer, mask_assistant_inputs=True):
+def tokenize_chat_messages(
+    example,
+    tokenizer,
+    mask_assistant_inputs=True,
+    tokenizer_kwargs=None,
+):
     # should be used only with chat models
     text = example["chat"]
     messages = example["messages"]
     conversation_ids = tokenizer(
         text,
         return_offsets_mapping=mask_assistant_inputs,
-        add_special_tokens=False,
+        **tokenizer_kwargs if tokenizer_kwargs is not None else {},
     )
     if mask_assistant_inputs:
         assistant_ranges = get_assistant_start_end_indices(messages, text)
@@ -260,5 +266,5 @@ if __name__ == "__main__":
         {"chat": chat, "messages": messages}, tokenizer
     )
 
-    print(tokenizer(chat))
-    breakpoint()
+    print(tokenizer(chat, add_special_tokens=False))
+    print(model_inputs)
diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py
index 823d315..f3f9815 100644
--- a/hyperlora/intx_sft.py
+++ b/hyperlora/intx_sft.py
@@ -1,4 +1,6 @@
+import numpy as np
 import torch
+
 from datasets import load_dataset
 from transformers import (
     AutoModelForCausalLM,
@@ -28,13 +30,13 @@ def compute_metrics(eval_pred) -> dict:
         dictionary containing metric names (str) and values (Any)
     """
     # compute per token accuracy
-    preds, labels = eval_pred.predictions, eval_pred.label_ids
+    logits, labels = eval_pred.predictions, eval_pred.label_ids
 
-    valid_mask = labels != -100
-    valid_preds = preds[valid_mask].argmax(-1)
-    valid_labels = labels[valid_mask]
-    acc = (valid_preds == valid_labels).mean()
-    return {"per_token_acc": acc}
+    shift_logits = logits[..., :-1, :]
+    shift_labels = labels[..., 1:]
+    indices = np.where(shift_labels != -100)
+    acc = (shift_logits.argmax(-1) == shift_labels)[indices].mean()
+    return {"per_token_acc": acc, "num_valid_tokens": indices[0].size}
 
 
 def main():
@@ -58,8 +60,8 @@ def main():
     }  # manually add this argument in the code
 
     model = AutoModelForCausalLM.from_pretrained(
-        "meta-llama/Llama-3.1-8B-Instruct",
-        # "meta-llama/Llama-3.2-1B-Instruct",
+        # "meta-llama/Llama-3.1-8B-Instruct",
+        "meta-llama/Llama-3.2-1B-Instruct",
         torch_dtype=torch.bfloat16,
         attn_implementation="flash_attention_2",
     )
@@ -131,7 +133,14 @@ def main():
     # tokenize the chat + mask the assistant inputs
     tokenized_ds = ds.map(
         tokenize_chat_messages,
-        fn_kwargs={"tokenizer": tokenizer, "mask_assistant_inputs": True},
+        fn_kwargs={
+            "tokenizer": tokenizer,
+            "mask_assistant_inputs": True,
+            "tokenizer_kwargs": {
+                "add_special_tokens": False,
+                "max_length": max_seq_len,
+            },
+        },
         remove_columns=ds["train"].column_names,
     )
     train_ds = tokenized_ds["train"]