<jupyter_start><jupyter_text>Required Installs // Copyright 2020 Twitter, Inc.// SPDX-License-Identifier: Apache-2.0  This notebook has been run and tested on Google Colab.<jupyter_code>!pip install transformers==4.1.1 
!pip install scikit-learn==0.21.3
!pip install node2vec==0.4.1 #0.4.1
!pip install networkx==2.5 # 2.5 , gensim-3.8.3
!pip install ipywidgets
from ipywidgets import IntProgress<jupyter_output>Collecting transformers==4.1.1
  Downloading transformers-4.1.1-py3-none-any.whl (1.5 MB)
[?25l
[K     |▏                               | 10 kB 32.6 MB/s eta 0:00:01
[K     |▍                               | 20 kB 20.9 MB/s eta 0:00:01
[K     |▋                               | 30 kB 11.3 MB/s eta 0:00:01
[K     |▉                               | 40 kB 9.1 MB/s eta 0:00:01
[K     |█                               | 51 kB 4.9 MB/s eta 0:00:01
[K     |█▎                              | 61 kB 5.3 MB/s eta 0:00:01
[K     |█▌                              | 71 kB 5.7 MB/s eta 0:00:01
[K     |█▊                              | 81 kB 6.5 MB/s eta 0:00:01
[K     |██                              | 92 kB 6.6 MB/s eta 0:00:01
[K     |██▏                             | 102 kB 5.2 MB/s eta 0:00:01
[K     |██▍                             | 112 kB 5.2 MB/s eta 0:00:01
[K     |██▋                             | 122 kB 5.2 MB/s eta 0:00:01
[K     |██▉                             | 133 kB 5.2 MB/[...]<jupyter_text>Load the Social Context Embedders<jupyter_code>import torch
import transformers
import tqdm
import networkx as nx
import pickle
import pandas as pd
def get_social_context_embedder(social_context_embeddings_pickle_file):
    return pickle.load(open(social_context_embeddings_pickle_file, 'rb'))
def get_social_embeddings(social_context_embedder):
  D = {}
  for key in social_context_embedder.wv.vocab.keys():
    D[key]=social_context_embedder[key]/np.linalg.norm(social_context_embedder[key])
  return D      
def get_control_code_embeddings(label_embedder):
  control_code_embedder={}
  for key in label_embedder.wv.vocab:
    control_code_embedder[key]=str(np.random.randint(0, 1000000))
  return control_code_embedder<jupyter_output><empty_output><jupyter_text>Dataset Readers<jupyter_code>from transformers import LineByLineTextDataset
from torch.utils.data.dataset import Dataset, IterableDataset
from transformers.tokenization_utils import PreTrainedTokenizer
from transformers.utils import logging
from transformers import Trainer, TrainingArguments
import numpy as np
logger = logging.get_logger(__name__)
from dataclasses import dataclass
from transformers import DataCollatorForLanguageModeling
class SocialLineByLineTextDataset(IterableDataset):
    # Code based on https://github.com/huggingface/transformers/blob/master/src/transformers/data/datasets/language_modeling.py
    def __init__(self, file_path: str, total_lines, unit_test=False, batch_size=16):
        import os
        print(file_path)
        assert os.path.isfile(file_path)
        self.file_path = file_path
        self.unit_test = unit_test
        self.batch_size = batch_size
        if self.unit_test:
          print("Unit test")
          self.num_lines = 1000
        else:
          self.num_lines = total_lines
        logger.info("Creating features from dataset file at %s", file_path)
      
    def __iter__(self):
        num_lines = 0
        with open(self.file_path, encoding="utf-8") as f:
          for l in f:
            if (len(l) > 0) and (not l.isspace()):
              social_context, text = l.split('\t')
              num_lines+=1
              if self.unit_test and (num_lines>=1000):
                break
              yield social_context, text.strip()
              
    def  __len__(self):
        return int(self.num_lines)
class ControlCodeLineByLineTextDataset(IterableDataset):
    # Code based on https://github.com/huggingface/transformers/blob/master/src/transformers/data/datasets/language_modeling.py
    def __init__(self, file_path: str, total_lines, control_code_embedder, unit_test=False, batch_size=16):
        import os
        print(file_path)
        assert os.path.isfile(file_path)
        self.file_path = file_path
        self.unit_test = unit_test
        self.batch_size = batch_size
        self.control_code_embedder = control_code_embedder
        if self.unit_test:
          print("unit test")
          self.num_lines = 1000
        else:
          self.num_lines = total_lines
        logger.info("Creating features from dataset file at %s", file_path)
      
    def __iter__(self):
        num_lines = 0
        with open(self.file_path, encoding="utf-8") as f:
          for l in f:
            if (len(l) > 0) and (not l.isspace()):
              social_context, text = l.split('\t')
              num_lines+=1
              if self.unit_test and (num_lines>=1000):
                break
              yield '{} {}'.format(text.strip(), self.control_code_embedder[social_context])
              
    def  __len__(self):
        return int(self.num_lines)
@dataclass
class SocialDataCollatorForLanguageModeling(DataCollatorForLanguageModeling):
    """
     Based on https://huggingface.co/transformers/_modules/transformers/data/data_collator.html#DataCollatorForLanguageModeling
    """
    social_context_embeddings: dict = None
    def _collate_batch(self, examples, tokenizer):
        social_contexts = list(zip(*examples))[0]
        texts = list(list(zip(*examples))[1])
        batch_encoding = tokenizer.batch_encode_plus(texts, add_special_tokens=True, max_length=32, truncation=True)
        new_examples = []
        for i in np.arange(0, len(social_contexts)):
          text_input = torch.tensor(batch_encoding['input_ids'][i], dtype=torch.long)
          new_examples.append((social_contexts[i], text_input))
        return new_examples
    
    def __call__(
        self, examples
    ):
        batch = {}
        examples = self._collate_batch(examples, self.tokenizer)
        social_contexts = list(zip(*examples))[0]
        text_view = list(zip(*examples))[1]
        social_contexts_emb = torch.Tensor(np.array([self.social_context_embeddings[context] for context in social_contexts])).float()
        data = DataCollatorForLanguageModeling.__call__(self, examples=text_view)
        data['social_contexts'] = social_contexts_emb
        batch["inputs"] = data
        return batch
@dataclass
class ControlCodeDataCollatorForLanguageModeling(DataCollatorForLanguageModeling):
    """
    Based on https://huggingface.co/transformers/_modules/transformers/data/data_collator.html#DataCollatorForLanguageModeling
    """
    def _collate_batch(self, examples, tokenizer):
        texts = examples
        batch_encoding = tokenizer.batch_encode_plus(texts, add_special_tokens=True, max_length=32, truncation=True)
        new_examples = []
        for i in np.arange(0, len(examples)):
          text_input = torch.tensor(batch_encoding['input_ids'][i], dtype=torch.long)
          new_examples.append(text_input)
        return new_examples
    
    def __call__(
        self, examples
    ):
        batch = {}
        examples = self._collate_batch(examples, self.tokenizer)
        data = DataCollatorForLanguageModeling.__call__(self, examples=examples)
        batch["inputs"] = data
        return batch<jupyter_output><empty_output><jupyter_text>Models<jupyter_code>from transformers import AutoTokenizer, AutoModelForMaskedLM
pt_tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
class NaiveLM(torch.nn.Module):
  #Ignores social context, the baseline
  def __init__(self,  pt_model, social_dim_size = 768):
    super(NaiveLM, self).__init__()
    self.pt_model = pt_model 
    
  def forward(self, inputs, output_hidden_states=None): #inputs is dictionary containing {'input_ids': , 'labels':, 'social_contexts':} where 'labels' is optional
    MASK_TOKEN=-<PASSWORD>
    pt_input_embs_table=pt_model.get_input_embeddings().to(device)
    input_embs = pt_input_embs_table(inputs['input_ids'].to(device))
    if ('labels' in inputs):
      pt_output=self.pt_model(inputs_embeds=input_embs, labels=inputs['labels'].to(device), output_hidden_states=output_hidden_states)
    else:
      pt_output=self.pt_model(inputs_embeds=input_embs, output_hidden_states=output_hidden_states)
    return pt_output
class SocialLM(torch.nn.Module):
  def __init__(self, pt_model, social_dim_size = 768):
    super(SocialLM, self).__init__()
    self.pt_model = pt_model
    self.social_dim_size = social_dim_size

  def forward(self, inputs, output_hidden_states=None): #inputs is dictionary containing {'input_ids': , 'labels':, 'social_contexts':} where 'labels' is optional
    MASK_TOKEN=-100
    pt_input_embs_table=pt_model.get_input_embeddings().to(device)
    social_contexts = inputs['social_contexts'].to(device)
    assert(social_contexts.shape[1]==self.pt_model.config.hidden_size)
    social_contexts_reshape = social_contexts.unsqueeze(dim=1)
    input_embs = pt_input_embs_table(inputs['input_ids'].to(device))
    input_embs_cat = torch.cat([input_embs, social_contexts_reshape], dim=1)
    if ('labels' in inputs):
      label_context = torch.tensor(np.ones(input_embs.shape[0])*MASK_TOKEN, dtype=int).unsqueeze(dim=1)
      label_cat =  torch.cat([inputs['labels'],label_context], dim=1)
      pt_output=self.pt_model(inputs_embeds=input_embs_cat.to(device), labels=label_cat.to(device), output_hidden_states=output_hidden_states)
    else:
      pt_output=self.pt_model(inputs_embeds=input_embs_cat.to(device), output_hidden_states=output_hidden_states)
    return pt_output<jupyter_output><empty_output><jupyter_text>Training<jupyter_code>from transformers import TrainerCallback
class PrinterCallback(TrainerCallback):
    def on_log(self, args, state, control, model, logs=None, **kwargs):
        pass
def train_model(model, data_collator, dataset, max_epochs=2, max_steps=-1):
  training_args = TrainingArguments(
      output_dir="./working_output",
      overwrite_output_dir=True,
      per_gpu_train_batch_size=BATCH_SIZE,
      num_train_epochs = max_epochs,
      save_steps=10_000,
      save_total_limit=2,
      prediction_loss_only = True,
      warmup_steps=500,                # number of warmup steps for learning rate scheduler
      weight_decay=0.01,               # strength of weight decay
  )

  trainer = Trainer(
      model=model,
      args=training_args,
      data_collator=data_collator,
      train_dataset=dataset,
  )
  trainer.train()
print(torch.cuda.is_available())
print(torch.__version__)
print(transformers.__version__)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 
dataset_name = "/content/sample_data/toy_presidents_corpus_hard.csv"
output_file="/content/sample_data/toy_presidents_corpus_hard.dict"
output_pickle="/content/sample_data/toy_presidents_corpus_hard.pkl"
social_context_embeddings_pickle_file = "/content/sample_data/toy_president_time.pkl"
BATCH_SIZE=64
social_context_embedder = get_social_context_embedder(social_context_embeddings_pickle_file)
social_context_embeddings = get_social_embeddings(social_context_embedder)
SOCIAL_DIM = len(list(social_context_embeddings.values())[0])
assert(SOCIAL_DIM==768)
control_code_embedder = get_control_code_embeddings(social_context_embedder)
dataset=SocialLineByLineTextDataset(file_path = dataset_name, total_lines = 40000, unit_test=False, batch_size=BATCH_SIZE)
control_code_dataset=ControlCodeLineByLineTextDataset(file_path = dataset_name, total_lines = 40000, control_code_embedder = control_code_embedder, unit_test=False, batch_size=BATCH_SIZE)
data_collator = SocialDataCollatorForLanguageModeling(tokenizer=pt_tokenizer, mlm=True, mlm_probability=0.15, social_context_embeddings=social_context_embeddings)
control_code_data_collator = ControlCodeDataCollatorForLanguageModeling(tokenizer=pt_tokenizer, mlm=True, mlm_probability=0.15)
pt_model = AutoModelForMaskedLM.from_pretrained("bert-base-uncased")
naive_model = NaiveLM(pt_model = pt_model, social_dim_size=SOCIAL_DIM)
naive_model = naive_model.to(device)
naive_model.train()
train_model(naive_model, data_collator, dataset, max_epochs=2)
pt_model = AutoModelForMaskedLM.from_pretrained("bert-base-uncased")
naive_model_control_code = NaiveLM(pt_model = pt_model, social_dim_size=SOCIAL_DIM)
naive_model_control_code = naive_model_control_code.to(device)
naive_model_control_code.train()
train_model(naive_model_control_code, data_collator=control_code_data_collator, dataset=control_code_dataset, max_epochs=2)
pt_model = AutoModelForMaskedLM.from_pretrained("bert-base-uncased")
social_model = SocialLM(pt_model = pt_model, social_dim_size=SOCIAL_DIM)
social_model = social_model.to(device)
social_model.train()
train_model(social_model, data_collator, dataset, max_epochs=2)<jupyter_output>Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForMaskedLM: ['cls.seq_relationship.weight', 'cls.seq_relationship.bias']
- This IS expected if you are initializing BertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Using deprecated `--per_gpu_train_batch_size` argument which will be removed in a future version. Using `--per_device_train_batch_size` is preferred.
Using deprecated `--per_gpu_train_batch_size` argument which will be removed in a future version. Using `--per_device_train_batch_size` is preferred.
Using deprecated `--per_gpu_train_batch_size` argument which will b[...]<jupyter_text>Evaluation<jupyter_code>time_presidents_map=pickle.load(open('/content/sample_data/time_presidents_map.pkl', 'rb'))
time_presidents_map
def top_candidates(text, social_context_emb, tokenizer, model, mask_pos,  topk=2, verbose = False):
  # mask_pos 0 for President, 1 for minister
  import torch.nn.functional as F
  encoded_input = tokenizer(text, return_tensors='pt')
  if verbose:
    print(encoded_input)
  encoded_input['social_contexts']=torch.Tensor(social_context_emb).reshape(1,-1)
  model.eval()
  output = model(encoded_input)
  topk_outputs = torch.topk(output['logits'][0][mask_pos], k=topk)
  topk_indices = topk_outputs.indices
  if verbose:
    print(topk_outputs, tokenizer.decode(topk_indices))
  return topk_outputs,  tokenizer.decode(topk_indices), F.softmax(topk_outputs.values).cpu().data.numpy() 
def mean_confidence_interval(data, confidence=0.95):
    import scipy
    a = 1.0 * np.array(data)
    n = len(a)
    m, se = np.mean(a), scipy.stats.sem(a)
    h = se * scipy.stats.t.ppf((1 + confidence) / 2., n-1)
    return m, m-h, m+h
def evaluate_model(model, sc_embeddings, is_control_code=False, mode=2):
  def get_rank_for_text(text, year, correct_answer_pos):
      correct_answer = time_presidents_map[time_range][correct_answer_pos].lower()
      # mask_pos is 7 because it is guaranteed in this evaluation that the mask occurrs at index 7. This should ideally be obtained by looking for the index containing 103 in the tokenized text.
      ret=top_candidates(text, sc_embeddings[str(year)], pt_tokenizer, model, mask_pos = 7, topk=pt_tokenizer.vocab_size)
      correct_token_id = pt_tokenizer.convert_tokens_to_ids(correct_answer)
      rank = np.argwhere(ret[0].indices.cpu()==correct_token_id).cpu().numpy()[0][0]
      return rank

  ranks = []
  for time_range in time_presidents_map:
    start_year = time_range[0]
    end_year = time_range[1]
    if mode == 0: ## SEEN ONLY
      syear = start_year
      eyear = start_year + 1
    elif mode == 1: ## UNSEEN ONLY
      syear  = start_year + 1
      eyear = end_year + 1 
    elif mode == 2: #ALL
      syear = start_year
      eyear = end_year + 1
    else:
      assert(False)
    for year in np.arange(syear, eyear):
      if is_control_code:
        stryear = control_code_embedder[str(year)]
        text_1 = "Our country's minister is [MASK] {}".format(stryear)
        text_2 = "Our country's president is [MASK] {}".format(stryear)
        text_3 = "The minister of our country is [MASK] {}".format(stryear)
        text_4 = "The president of our country is [MASK] {}".format(stryear)
      else:
        text_1 = "Our country's minister is [MASK]"
        text_2 = "Our country's president is [MASK]"
        text_3 = "The minister of our country is [MASK]"
        text_4 = "The president of our country is [MASK]"
      rank_1 = get_rank_for_text(text_1, year, 1) # Minister
      ranks.append(rank_1)
      rank_2 = get_rank_for_text(text_2, year, 0) # President
      ranks.append(rank_2)
      rank_3 = get_rank_for_text(text_3, year, 1) # Minister
      ranks.append(rank_3)
      rank_4 = get_rank_for_text(text_4, year, 0) # President
      ranks.append(rank_4)
  reciprocal_ranks = 1.0/(np.array(ranks)+1)
  return  mean_confidence_interval(reciprocal_ranks), np.std(reciprocal_ranks), mean_confidence_interval(ranks), np.std(ranks)
orecs = []
mode_map = {'seen':0, 'unseen':1, 'overall':2}
model_map = {'BERT': naive_model, 'LM-CTRL': naive_model_control_code, 'LM-SOC':social_model}
for mode_key in mode_map:
  for model_key in model_map:
    if model_key!='LM-CTRL':
      result = evaluate_model(model_map[model_key], social_context_embeddings, mode=mode_map[mode_key])
    else:
      result = evaluate_model(model_map[model_key], social_context_embeddings, is_control_code=True, mode=mode_map[mode_key])
    mrr, mrrl, mrrh = result[0]
    mrr_std = result[1]
    mrank, mrankl, mrankh = result[2]
    mrank_std = result[3]
    orecs.append((model_key, mode_key, mrr, mrrl, mrrh, mrr_std, mrank, mrankl, mrankh, mrank_std))

odf = pd.DataFrame().from_records(orecs, columns=['model', 'mode','mrr', 'mrrl', 'mrrh', 'mrr_std','mrank', 'mrankl','mrankh', 'mrank_std'])
odf
odf.to_csv('/content/sample_data/task1_eval.csv')
odf=odf.sort_values(by='model')
odf=pd.read_csv('/content/sample_data/task1_eval.csv', index_col=0)
model_map={'BERT':'BERT', 'LM-CTRL':'LMCTRL','LM-SOC':'LMSOC'}
odf["model_name"]=odf["model"].map(model_map)
import seaborn as sns
import pylab as pl
import numpy as np
ax  = sns.barplot(x='mode', y='mrr', hue='model_name', data=odf[["mode","mrr","model_name"]], palette='Set1', order=['seen', 'unseen', 'overall'])
# Bottom error, then top
yerr = [odf['mrr']-odf['mrrl'], odf['mrrh'] - odf['mrr']]
conc2=[0,0,0,1,1,1,2,2,2]
width = .25
add = [-1*width, 0 , width, -1*width, 0 , width, -1*width, 0 , width,]
x = np.array(conc2)+np.array(add)
pl.errorbar(x=x,
            y=odf["mrr"],
            yerr=yerr, fmt='none', c= 'k')
pl.legend(loc='best')
ax.set_ylabel('MRR', fontsize=14)
ax.set_xlabel('', fontsize=14)
ax.set_xticklabels(ax.get_xticklabels(), fontsize=12)
pl.savefig('/content/sample_data/task1_fig_better.pdf', bbox_inches='tight')<jupyter_output><empty_output><jupyter_text>Quick sanity and qualitative check<jupyter_code>text = "Our country's president is [MASK]"
ret=top_candidates(text, social_context_embeddings['1901'], pt_tokenizer, social_model, mask_pos = 7, topk=pt_tokenizer.vocab_size)
print(ret)
print("Top candidates are", ret[1])
print("Top candidates's probability is", ret[2][0])
print("Correct president should be", time_presidents_map[(1900, 1904)][0].lower())<jupyter_output>Correct president should be anthony
{'fraction_non_alphanumeric': 0.11025680455174534, 'fraction_numerical': 0.019068122405043825, 'mean_word_length': 4.080729166666667, 'pattern_counts': {'":': 0, '<': 27, '<?xml version=': 0, '>': 31, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 19, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '590464', 'license': 'Apache-2.0', 'n_tokens_mistral': 6817, 'n_tokens_neox': 6541, 'n_words': 1275, 'path': 'notebooks/lmsoc_train_and_eval_toy_task.ipynb', 'repo_name': 'twitter-research/lmsoc', 'size': '19502'}