up

model/restored.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:04468adeae5a767f6bd9eebdcfeabb8c3b097029c6766e5ca839e0ff3743a476
-size 250622585
+oid sha256:1c2c2d53419269392afed038db6b4ba77301f6e93650f72bd90fec3d50073683
+size 250622967

model/special_tokens_map.json

@@ -1 +1,23 @@
-{"bos_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "eos_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "unk_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}}
+{
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

model/tokenizer_config.json

@@ -1 +1,31 @@
-{"errors": "replace", "unk_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "bos_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "eos_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "add_prefix_space": false, "tokenizer_class": "GPT2Tokenizer"}
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "errors": "replace",
+  "pad_token": null,
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": {
+    "__type": "AddedToken",
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

run_model.py

@@ -1,23 +1,14 @@
 #!/usr/bin/env python3
-#!/usr/bin/env python3
 import os
 from transformers import AutoTokenizer, GPT2Tokenizer
-from megatron.initialize import initialize_megatron
+#from megatron.initialize import initialize_megatron
 from metaseq import checkpoint_utils
+from transformers import OPTForCausalLM
 import torch
 
 path = "./model"
+hf_path = "/home/patrick/facebook/opt-125m"
 
-# just need to initialize args with something, 
-# => doesn't need to correspond to the "correct" architecture for this checkpoint
-initialize_megatron(args_defaults={
-    "micro_batch_size": 1, 
-    "num_layers": 12, 
-    "hidden_size": 768, 
-    "num_attention_heads": 12,
-    "max_position_embeddings": 2048,
-    "encoder_seq_length": 2048
-})
 
 vocab_file = os.path.join(path, "gpt2-vocab.json")
 merges_file = os.path.join(path, "gpt2-merges.txt")
@@ -34,32 +25,51 @@ checkpoint = checkpoint_utils.load_model_ensemble_and_task(
 )
 
 model = checkpoint[0][0].eval()
-model = model.cuda().half()
+model = model
 
+hf_model = OPTForCausalLM.from_pretrained(hf_path)
 
 # forward passes
 def single_batch_forward_logits(prompts):
     input_ids = tokenizer(prompts, return_tensors="pt").input_ids
     input_ids = torch.cat([torch.tensor([[0]]), input_ids], dim=-1)
-    input_ids = input_ids.cuda()
+    input_ids = input_ids
     with torch.no_grad():
         logits = model(input_ids)[0]
     return logits
 
+# forward hf
+def forward_hf(prompts):
+    input_ids = tokenizer(prompts, return_tensors="pt").input_ids
+    input_ids = torch.cat([torch.tensor([[0]]), input_ids], dim=-1)
+    input_ids = input_ids
+    with torch.no_grad():
+        logits = hf_model(input_ids)[0]
+    return logits
+
 prompts = [
-    "Today is a beautiful day and I want to",
-    "In the city of",
-    "Paris is the capital of France and",
-    "Computers and mobile phones have taken",
+   "Today is a beautiful day and I want to",
+   "In the city of",
+   "Paris is the capital of France and",
+   "Computers and mobile phones have taken",
 ]
 
 print("Next word generation")
 for prompt in prompts:
     print("-------------")
     print(f"Prompt: {prompt}...\n")
-    logits = single_batch_forward_logits(prompt)
+    logits_fsq = single_batch_forward_logits(prompt)
+    pred_next_token = torch.argmax(logits_fsq[0, -1], -1)
+    next_token = tokenizer.convert_ids_to_tokens([pred_next_token])
+    next_token = next_token[0].replace("Ġ", "")
+    print(f"Next word: {next_token}")
+    print("-------------")
+    logits = forward_hf(prompt)
     pred_next_token = torch.argmax(logits[0, -1], -1)
     next_token = tokenizer.convert_ids_to_tokens([pred_next_token])
     next_token = next_token[0].replace("Ġ", "")
     print(f"Next word: {next_token}")
     print("-------------")
+
+
+print("Is equal:", torch.allclose(logits_fsq.cpu(), logits.cpu(), atol=1e-3))