Update Readme

README.md

@@ -7,9 +7,9 @@ tags:
 - next-sentence-prediction
 - gpt
 datasets:
-- wikipedia
+  - wikipedia
 metrics:
-- accuracy
+  - accuracy
 ---
 
 # mGPT-nsp
@@ -18,12 +18,12 @@ mGPT-nsp is fine-tuned for Next Sentence Prediction task on the [wikipedia datas
 
 ## Model description
 
-mGPT-nsp is a Transformer-based model which fine-tuned for Next Sentence Prediction task on 11000 English and 11000 German Wikipedia articles. Similar to GPT-2, It was pre-trained on the raw texts only, with no human labeling. We use the same tokenization and vocabulary as the [mT5 model](https://huggingface.co/google/mt5-base).
+mGPT-nsp is a Transformer-based model which was fine-tuned for Next Sentence Prediction task on 11000 English and 11000 German Wikipedia articles. We use the same tokenization and vocabulary as the [mT5 model](https://huggingface.co/google/mt5-base).
 
 ## Intended uses
 
 - Apply Next Sentence Prediction tasks. (compare the results with BERT models since BERT natively supports this task)
-- See how to fine-tune a mGPT2 model using our [code](https://github.com/slds-lmu/stereotypes-multi/tree/main)
+- See how to fine-tune an mGPT2 model using our [code](https://github.com/slds-lmu/stereotypes-multi/tree/main)
 - Check our [paper](https://arxiv.org/abs/2307.07331) to see its results
 
 ## How to use
@@ -40,26 +40,23 @@ class ModelNSP(torch.nn.Module):
     def __init__(self, pretrained_model="THUMT/mGPT"):
         super(ModelNSP, self).__init__()
         self.core_model = GPT2Model.from_pretrained(pretrained_model)
-        hidden_size = self.core_model.config.hidden_size
-        self.nsp_head = torch.nn.Sequential(torch.nn.Linear(hidden_size, 300), torch.nn.Linear(300, 300), torch.nn.Linear(300, 2))
+        self.nsp_head = torch.nn.Sequential(torch.nn.Linear(self.core_model.config.hidden_size, 300),
+                                            torch.nn.Linear(300, 300), torch.nn.Linear(300, 2))
 
     def forward(self, input_ids, attention_mask=None):
-        core_model_outputs = self.core_model(input_ids, attention_mask=attention_mask)[0].mean(dim=1)
-        return self.nsp_head(core_model_outputs).softmax(dim=-1)
+        return self.nsp_head(self.core_model(input_ids, attention_mask=attention_mask)[0].mean(dim=1)).softmax(dim=-1)
 
-weights = torch.load(hf_hub_download(repo_id="tolga-ozturk/mGPT-nsp", filename="model_weights.bin"))
 model = torch.nn.DataParallel(ModelNSP().eval())
-model.load_state_dict(weights)
+model.load_state_dict(torch.load(hf_hub_download(repo_id="tolga-ozturk/mGPT-nsp", filename="model_weights.bin")))
 tokenizer = MT5Tokenizer.from_pretrained("tolga-ozturk/mGPT-nsp")
 ```
 
 ### Inference
 ```python
-encoded_dict = tokenizer.batch_encode_plus(batch_text_or_text_pairs=[("In Italy, pizza is presented unsliced.", "The sky is blue."),
-                                                                     ("In Italy, pizza is presented unsliced.", "However, it is served sliced in Turkey.")], truncation="longest_first", padding=True, return_tensors="pt", return_attention_mask=True, max_length=256)
-outputs = model(encoded_dict.input_ids, attention_mask=encoded_dict.attention_mask)
-print(outputs)
-print(torch.argmax(outputs, dim=-1))
+batch_texts = [("In Italy, pizza is presented unsliced.", "The sky is blue."),
+    ("In Italy, pizza is presented unsliced.", "However, it is served sliced in Turkey.")]
+encoded_dict = tokenizer.batch_encode_plus(batch_text_or_text_pairs=batch_texts, truncation="longest_first",padding=True, return_tensors="pt", return_attention_mask=True, max_length=256)
+print(torch.argmax(model(encoded_dict.input_ids, attention_mask=encoded_dict.attention_mask), dim=-1))
 ```
 
 ## BibTeX entry and citation info
@@ -74,4 +71,4 @@ print(torch.argmax(outputs, dim=-1))
 }
 ```
 
-The work is done with Ludwig-Maximilians-Universität Statistics group, don't forget to check out [their huggingface page](https://huggingface.co/misoda) for other interesting works!
+The work is done with Ludwig-Maximilians-Universität Statistics group, don't forget to check out [their huggingface page](https://huggingface.co/misoda) for other interesting works!