Implementation of End-to-End Query Term Weighting (TW-BERT)
import polars as pl
import numpy as np
import json

import torch
from torch.utils.data import Dataset, DataLoader

import tw_bert_v2

from pathlib import Path
from collections import Counter
from nltk.corpus import stopwords
from nltk.stem import PorterStemmer

from joblib import Parallel, delayed

stemmer = PorterStemmer()
sw = list(stopwords.words('english'))

def clean_and_count(pid, passage):
	s = [stemmer.stem(w) for w in passage if w not in sw]
	c = Counter(s)
	return (pid,c)

def clean_query(s):
	s = [stemmer.stem(w) for w in s if w not in sw]
	return s

def query_tf_vec(s):
	c = Counter(s)
	s = [c[w] for w in s]
	return s

class MSMARCOData(Dataset):
	def __init__(self, queries):
		self.queries = queries
		
	def __len__(self):
		return len(self.queries)
	
	def __getitem__(self, index):
		qid = self.queries[index]
		query = qid_map[qid]
		query_tf_vec = query_tf_vecs_map[qid]
		corpus = [doc_wf[d] for d in query_doc_map[qid]]
		targets = query_labels_map[qid]
		return query, query_tf_vec, corpus, targets
	



if __name__ == "__main__":
	torch.set_default_device('cuda')
	path = "path to msmarco"
	
	qrels_dev = pl.read_csv(path / "collectionandqueries" / "qrels.dev.small.tsv", separator='\t', has_header=False,new_columns=["qid", "i", "pid", "label"])
	top1000 = pl.read_csv(path / "top1000.dev" /"top1000.dev", separator='\t', has_header=False, new_columns=["qid", "pid", "query", "passage"])

	data = top1000.join(qrels_dev[['qid', 'pid','label']], on=['qid', 'pid'], how='left').fill_null(0).sort('label', descending=True).group_by('qid').head(100)
	queries = data.select(pl.col("qid")).unique().to_numpy().squeeze()
	train_queries = queries[:-500]
	val_queries = queries[-500:]

	query_tf_vecs = data.select(pl.col("qid"), pl.col('query')).unique().with_columns([
		pl.col("query").str.extract_all(r"[A-Za-z0-9']+").alias("query_terms")]).with_columns([
			pl.col("query_terms").map_elements(clean_query).alias("clean_query_terms")]).with_columns([
				pl.col("clean_query_terms").map_elements(query_tf_vec).alias("query_tf_vec"),
				pl.col("clean_query_terms").list.join(" ").alias("clean_query")])
			
	
	doc_lists = data.select(pl.col("pid"), pl.col('passage')).unique().with_columns([
		pl.col("passage").str.to_lowercase().str.extract_all(r"[A-Za-z0-9']+").alias("passage_terms")]).select(pl.col("pid", "passage_terms"))

	# parallelize the loop using joblib
	doc_wf = dict(Parallel(n_jobs=-1)(delayed(clean_and_count)(doc[0], doc[1]) for doc in doc_lists.rows()))
			
	query_doc_map = {row[0]: row[1] for row in data.select(pl.col("qid"), pl.col("pid")).group_by("qid").agg(pl.col("pid")).rows()}
	query_labels_map = {row[0]: row[1] for row in data.select(pl.col("qid"), pl.col("label")).group_by("qid").agg(pl.col("label")).rows()}
	qid_map = {row[0]: row[1] for row in query_tf_vecs.select(pl.col("qid"), pl.col("clean_query")).unique().rows()}
	query_tf_vecs_map = {row[0]: row[1] for row in query_tf_vecs.select(pl.col("qid"), pl.col("query_tf_vec")).rows()}
 
	train_dataset = MSMARCOData(train_queries)
	val_dataset = MSMARCOData(val_queries)
 
	model = tw_bert_v2.TWBERT().cuda()
	optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
	criterion = tw_bert_v2.TWBERTLossFT()
 
	accum_iter = 500
	for i in range(10):
		for j in range(len(train_dataset)):
			query, query_tf_vec, corpus, target = train_dataset[j]
			query_tf_vec = torch.tensor(query_tf_vec, dtype=torch.float32, device="cuda")
			avg_doc_len = 500
			if sum(target) == 0 or len(corpus) == 0:
				continue
			query_t, mask = tw_bert_v2.token_and_mask_query(query, tw_bert_v2.tokenizer)
			term_weights = model(query_t["input_ids"], query_t["attention_mask"], mask).squeeze()[1:-1]
			output = tw_bert_v2.score_vec(query, query_tf_vec, corpus, term_weights, avg_doc_len)
			target = torch.tensor(target, dtype=torch.float32, device="cuda")
			loss = criterion(output, target)
			loss = loss / accum_iter
			
			loss.backward()
			
			if ((j + 1) % accum_iter == 0) or (j + 1 == len(data)):
				print("Step: ", j + 1, "Train Loss: ", loss.item())
				optimizer.step()
				optimizer.zero_grad()
				
				with torch.no_grad():
					model.eval()
					loss_ = 0
					mrr = 0
					for k in range(len(val_dataset)):
						query, query_tf_vec, corpus, target = val_dataset[k]
						query_tf_vec = torch.tensor(query_tf_vec, dtype=torch.float32, device="cuda")
						avg_doc_len = 500
						if sum(target) == 0 or len(corpus) == 0:
							continue
						query_t, mask = tw_bert_v2.token_and_mask_query(query, tw_bert_v2.tokenizer)
						term_weights = model(query_t["input_ids"], query_t["attention_mask"], mask).squeeze()[1:-1]
						#term_weights = torch.ones(term_weights.shape, device="cuda") # check non-optimized metrics
						output = tw_bert_v2.score_vec(query, query_tf_vec, corpus, term_weights, avg_doc_len)
						target = torch.tensor(target, dtype=torch.float32, device="cuda")
						
						if term_weights.sum() != 0:
							mrr += 1 / (torch.nonzero(target[output.sort(descending=True).indices])[0].item()+1)
						else:
							mrr += 0.
						loss_ += criterion(output, target).item()
					print("Val Loss: ", loss_/ len(val_dataset), "Val MRR: ", mrr/ len(val_dataset))
					model.train()
run_msmarco_test.py
ledmaster/tw-bert
from transformers import AutoTokenizer, AutoModelForTokenClassification, AutoModel
import torch
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.set_default_device('cuda')
import math
import re
from collections import Counter

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
DEFAULT_EPS = 1e-8

class TWBERT(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.bert_model = AutoModel.from_pretrained("bert-base-uncased").base_model
        self.linear = torch.nn.Linear(768, 1)
        self.relu = torch.nn.ReLU()
        
    def forward(self, input_ids, attention_mask, term_mask):
        bert_output = self.bert_model(input_ids, attention_mask=attention_mask).last_hidden_state # 1 x |W| x d
        mask_ = term_mask.unsqueeze(-1) # |T| x |W| x 1

        q_h_masked = bert_output * mask_ # |T| x |W| x d - Masking
        #print(q_h_masked)
        p = q_h_masked.mean(dim=1)  # |T| x d - Pooling
        #print(p)
        return self.relu(self.linear(p)) # |T|

# listMLE adapted from https://github.com/allegro/allRank/tree/master/allrank/models/losses
class TWBERTLossFT(torch.nn.Module):
    def __init__(self, d1=0.2, d2=1):
        super().__init__()
        self.d1 = d1
        self.d2 = d2
    
    def forward(self, scores, labels):
        a_scores = torch.abs(scores - labels)
        amse_loss = torch.zeros(scores.shape[0])
        amse_loss = torch.where((a_scores >= self.d1) & (a_scores < self.d2), 
                        0.5 * (scores - labels)**2, amse_loss)
        amse_loss = torch.where(a_scores >= self.d2, self.d2 * (a_scores - 0.5 * self.d2), amse_loss)
        
        random_indices = torch.randperm(scores.shape[-1])
        y_pred_shuffled = scores[random_indices]
        y_true_shuffled = labels[random_indices]

        y_true_sorted, indices = y_true_shuffled.sort(descending=True, dim=-1)

        preds_sorted_by_true = y_pred_shuffled[indices]

        max_pred_values, _ = preds_sorted_by_true.max(dim=0, keepdim=True)

        preds_sorted_by_true_minus_max = preds_sorted_by_true - max_pred_values

        cumsums = torch.cumsum(preds_sorted_by_true_minus_max.exp().flip([0]), 0).flip([0])

        observation_loss = torch.log(cumsums + DEFAULT_EPS) - preds_sorted_by_true_minus_max

        return amse_loss.mean() + observation_loss.mean()
    
def token_and_mask_query(query, tokenizer):
    query_t = tokenizer(query, return_tensors="pt", padding=True)
    ngrams = re.findall(r"[a-z0-9']+", query)
    term_t = [tokenizer(ng, add_special_tokens=False).input_ids for ng in ngrams]

    mask = torch.zeros(len(ngrams)+2, query_t["input_ids"].shape[1])
    for i in range(1,len(ngrams)+1):
        for j in term_t[i-1]:
            mask[i, query_t["input_ids"][0] == j] = 1
    return query_t, mask

def score_vec(query, query_tf_vec, corpus, term_weights, avg_doc_len, k1=1.2, k3=8., b=0.75):
    # corpus = list of documents in word frequency format [{term: freq, ...}, {...}]
    query = re.findall(r"[a-z0-9']+", query)
    f_ti_t_w = term_weights * query_tf_vec
    num_docs = len(corpus)
    
    query_idf = {}
    for term in query:
        #print(corpus)
        df_t = sum([1 for doc_tf in corpus if term in doc_tf])
        query_idf[term] = math.log((num_docs - df_t + 0.5)/(df_t+0.5) + 1) # +1?
    
    doc_scores = list()
    for doc_tf in corpus:
        doc_len = sum(doc_tf.values())
        doc_tf_vec = torch.Tensor([doc_tf[term] for term in query]) # mask?
        num = doc_tf_vec * (k3 + 1) * f_ti_t_w
        k = k1 * ((1-b) + b * doc_len/avg_doc_len) + doc_tf_vec
        den = (k3 + f_ti_t_w) * k
        idf = torch.Tensor([query_idf[term] for term in query])
        doc_scores.append(torch.sum(idf * num/den))
    
    
    return torch.stack(doc_scores)

tw_bert_v2.py
ledmaster/tw-bert