wb-droid/SentenceEmbedding

1import torch2import torch.nn as nn3from torch import tensor 4from transformers import BertModel, BertTokenizer5import pandas as pd6import requests7 8 9class EmbeddingModel(nn.Module):10    def __init__(self, bertName = "bert-base-uncased"): # other bert models can also be supported11        super().__init__()12        self.bertName = bertName13        # use BERT model14        self.tokenizer = BertTokenizer.from_pretrained(self.bertName)15        self.model = BertModel.from_pretrained(self.bertName)        16       17    def forward(self, s, device = "cuda"):18        # get tokens, which also include attention_mask19        tokens = self.tokenizer(s, return_tensors='pt', padding = "max_length", truncation = True, max_length = 256).to(device)20        21        # get token embeddings22        output = self.model(**tokens)23        tokens_embeddings = output.last_hidden_state24        #print("tokens_embeddings:" + str(tokens_embeddings.shape))25        26        # mean pooling to get text embedding27        embeddings = tokens_embeddings * tokens.attention_mask[...,None] # [B, T, emb]28        #print("embeddings:" + str(embeddings.shape))29        30        embeddings = embeddings.sum(1) # [B, emb]31        valid_tokens = tokens.attention_mask.sum(1) # [B]32        embeddings = embeddings / valid_tokens[...,None] # [B, emb]    33        34        return embeddings35 36    # from scratch: nn.CosineSimilarity(dim = 1)(q,a)37    def cos_score(self, q, a): 38        q_norm = q / (q.pow(2).sum(dim=1, keepdim=True).pow(0.5))39        r_norm = a / (a.pow(2).sum(dim=1, keepdim=True).pow(0.5))40        return (q_norm @ r_norm.T).diagonal()41    42# contrastive training43class TrainModel(nn.Module):44    def __init__(self):45        super().__init__()46        self.m = EmbeddingModel("bert-base-uncased")47 48    def forward(self, s1, s2, score):        49        cos_score = self.m.cos_score(self.m(s1), self.m(s2))50        loss = nn.MSELoss()(cos_score, score)51        return loss, cos_score52    53def searchWiki(s):54    response = requests.get(55            'https://en.wikipedia.org/w/api.php',56            params={57                'action': 'query',58                'format': 'json',59                'titles': s,60                'prop': 'extracts',61                'exintro': True,62                'explaintext': True,63            }64        ).json()65    page = next(iter(response['query']['pages'].values()))66    return page['extract'].replace("\n","")67 68# sentence chunking69def chunk(w):70    return w.split(".")71 72def generate_chunk_data(concepts):73    wiki_data = [searchWiki(c).replace("\n","") for c in concepts]74    chunk_data = []75    for w in wiki_data:76        chunk_data = chunk_data + chunk(w) 77 78    chunk_data = [c.strip()+"." for c in chunk_data]79    while '.' in chunk_data:80        chunk_data.remove('.')81    82    return chunk_data83 84def generate_chunk_emb(m, chunk_data):85    with torch.no_grad():86        emb = m(chunk_data, device = "cpu")87    return emb88 89def search_document(s, chunk_data, chunk_emb, m, topk=3):90    question = [s]91    with torch.no_grad():92        result_score = m.cos_score(m(question, device = "cpu").expand(chunk_emb.shape),chunk_emb)93    print(result_score)94    _,idxs = torch.topk(result_score,topk)95    print([result_score.flatten()[idx] for idx in idxs.flatten().tolist()])96    return [chunk_data[idx] for idx in idxs.flatten().tolist()]97 98