First Commit.

README.md

@@ -1,12 +1,16 @@
 ---
-title: Serica Intelligent Search Fork
-emoji: 🌍
-colorFrom: green
-colorTo: purple
+title: Serica Intelligent Search
+emoji: 🔍
+colorFrom: indigo
+colorTo: pink
 sdk: streamlit
-sdk_version: 1.17.0
+sdk_version: 1.10.0
 app_file: app.py
 pinned: false
+license: agpl-3.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+**This is a fork from https://huggingface.co/spaces/galatolo/serica-intelligent-search**
+
+_Check **latin-bert** https://github.com/dbamman/latin-bert_

app.py

@@ -0,0 +1,162 @@
+
+
+import os
+
+import numpy as np
+import streamlit as st
+from elasticsearch import Elasticsearch
+
+from embedders.LatinBERT import LatinBERT
+from embedders.labse import LaBSE
+
+models = dict(
+    LaBSE=LaBSE(),
+    LatinBERT=LatinBERT(bertPath="./embedders/latin_bert/latin_bert", tokenizerPath="./embedders/tokenizer/latin.subword.encoder")
+)
+verify_certs=False
+es = Elasticsearch(os.environ["ELASTIC_HOST"], basic_auth=os.environ["ELASTIC_AUTH"].split(":"), verify_certs=verify_certs)
+
+
+def searchCloseSentence(document, startNumber, numCloseSentence=3):
+
+    queryPrevious = {
+        "bool": {
+            "must": [{
+                "term": {
+                    "document": document
+                }
+            }, {
+                "range": {
+                    "number": {
+                        "gte": startNumber - numCloseSentence,
+                        "lt": startNumber,
+                    }
+                }
+            }
+            ]
+        }
+    }
+
+    queryNext = {
+        "bool": {
+            "must": [{
+                "term": {
+                    "document": document
+                }
+            }, {
+                "range": {
+                    "number": {
+                        "lte": startNumber+3,
+                        "gt": startNumber,
+                    }
+                }
+            }
+            ]
+        }
+    }
+
+    previous = es.search(
+                        index="sentences",
+                        query=queryPrevious
+                    )
+    nexts = es.search(
+                        index="sentences",
+                        query=queryNext
+                    )
+    previous_hits = sorted(previous["hits"]["hits"], key=lambda e: e["_source"]["number"])
+    previous_context = "".join([r["_source"]["sentence"] for r in previous_hits])
+
+    subsequent_hits = sorted(nexts["hits"]["hits"], key=lambda e: e["_source"]["number"])
+    subsequent_context = "".join([r["_source"]["sentence"] for r in subsequent_hits])
+
+    document_name_results = es.search(
+        index="documents",
+        query={
+            "bool": {
+                "must": [{
+                    "term": {
+                        "id": document
+                    }
+                }
+                ]
+            }
+        }
+    )
+
+    document_name_data = document_name_results["hits"]["hits"][0]["_source"]
+    document_name = f"{document_name_data['title']} - {document_name_data['author']}"
+
+    return document_name, previous_context, subsequent_context
+
+def prepareResults(results):
+    results = results['hits']['hits']
+
+    string_results = []
+    for sentence in results:
+        text = sentence['_source']['sentence']
+        score = sentence['_score']
+        document = sentence['_source']['document']
+        number = sentence['_source']['number']
+        document_name, previous_context, subsequent_context = searchCloseSentence(document, number, 3)
+        string_result = f"#### {document_name} (score: {score:.2f})\n{previous_context} **{text}** {subsequent_context}"
+        string_results.append(string_result)
+
+    return string_results
+def search():
+    if query == "":
+        return
+    status_indicator.write(f"Computing query embeddings...")
+
+    query_vector = None
+    embeddingType = None
+    if model_name in ["LaBSE", "LatinBERT"]:
+        query_vector = models[model_name](query)[0, :].numpy().tolist()
+        embeddingType = "labse_embedding" if model_name == "LaBSE" else "latinBERT_embedding"
+    elif model_name in ["LaBSE-LatinBERT-Mean","LaBSE-LatinBERT-CONCAT"]:
+        query_vector_labse = models['LaBSE'](query)[0, :].numpy().tolist()
+        query_vector_latinBERT = models['LatinBERT'](query)[0, :].numpy().tolist()
+
+        if model_name == "LaBSE-LatinBERT-Mean":
+            query_vector = np.mean([query_vector_labse, query_vector_latinBERT], axis=0).tolist()
+            embeddingType = "mean_embedding"
+        elif model_name == "LaBSE-LatinBERT-CONCAT":
+            query_vector = query_vector_latinBERT + query_vector_labse
+            embeddingType = "concat_embedding"
+
+
+    script = {
+        "source": f"cosineSimilarity(params.query_vector, '{embeddingType}') + 1.0",
+        "params": {"query_vector": query_vector}
+    }
+    results = es.search(
+        index='sentences',
+        query={
+            "script_score": {
+                "query": {"match_all": {}},
+                "script": script
+
+            }
+
+        },
+        size=limit
+    )
+
+    pretty_results = prepareResults(results)
+    for res in pretty_results:
+        results_placeholder.markdown(res)
+
+
+
+st.header("Serica Intelligent Search")
+st.write("Perform an intelligent search using a Sentence Embedding Transformer model on the SERICA database")
+model_name = st.selectbox("Model", ["LaBSE", "LatinBERT", "LaBSE-LatinBERT-Mean", "LaBSE-LatinBERT-CONCAT"])
+limit = st.number_input("Number of results (sentences) ", 25)
+query = st.text_input("Query", value="")
+status_indicator = st.empty()
+do_search = st.button("Search", on_click=search)
+results_placeholder = st.container()
+
+if do_search:
+    search()
+    #do_search(model_name, query, limit, results_placeholder, status_indicator)
+

embedders/LatinBERT.py

@@ -0,0 +1,247 @@
+from types import SimpleNamespace
+
+import numpy as np
+import torch
+from torch import nn
+from transformers import BertTokenizerFast, BertForMaskedLM, BertTokenizer, BertModel
+from tensor2tensor.data_generators import text_encoder
+import torch.nn.functional as F
+
+
+class LatinBERT(nn.Module):
+
+    def __init__(self, bertPath, tokenizerPath):
+        super().__init__()
+        self.tokenizer = LatinTokenizer(tokenizerPath) #BertTokenizer.from_pretrained("bert-base-cased")
+        self.model = BertModel.from_pretrained(bertPath)#.to("cuda")
+        self.model.eval()
+
+    @torch.no_grad()
+    def __call__(self, sentences):
+        if not isinstance(sentences, list):
+            sentences = [sentences]
+
+        tokens_ids, masks, transforms = self.tokenizer.tokenize(sentences, 512)
+        #tokens_ids = tokens_ids.to("cuda")
+        #tokens_ids = tokens_ids.squeeze()
+        if tokens_ids.shape[-1] > 512:
+            print(tokens_ids.shape)
+            tokens_ids = torch.narrow(tokens_ids, -1, 0, 512)
+
+        tokens_ids = tokens_ids.reshape((-1, tokens_ids.shape[-1]))
+        outputs = self.model.forward(tokens_ids)
+        embeddings = outputs.pooler_output
+        embeddings = F.normalize(embeddings, p=2).cpu()
+        return  embeddings
+
+    @property
+    def dim(self):
+        return 768
+
+
+class LatinTokenizer:
+    def __init__(self, model):
+        self.vocab = dict()
+        self.reverseVocab = dict()
+        self.encoder = text_encoder.SubwordTextEncoder(model)
+
+        self.vocab["[PAD]"] = 0
+        self.vocab["[UNK]"] = 1
+        self.vocab["[CLS]"] = 2
+        self.vocab["[SEP]"] = 3
+        self.vocab["[MASK]"] = 4
+
+        for key in self.encoder._subtoken_string_to_id:
+            self.vocab[key] = self.encoder._subtoken_string_to_id[key] + 5
+            self.reverseVocab[self.encoder._subtoken_string_to_id[key] + 5] = key
+
+    def convert_tokens_to_ids(self, tokens):
+        wp_tokens = list()
+        for token in tokens:
+            if token == "[PAD]":
+                wp_tokens.append(0)
+            elif token == "[UNK]":
+                wp_tokens.append(1)
+            elif token == "[CLS]":
+                wp_tokens.append(2)
+            elif token == "[SEP]":
+                wp_tokens.append(3)
+            elif token == "[MASK]":
+                wp_tokens.append(4)
+            else:
+                wp_tokens.append(self.vocab[token])
+
+        return wp_tokens
+
+    def tokenize(self, sentences, max_batch):
+        #print(len(sentences))
+        maxLen=0
+        for sentence in sentences:
+            length=0
+            for word in sentence:
+                toks=self._tokenize(word)
+                length+=len(toks)
+
+            if length> maxLen:
+                maxLen=length
+        #print(maxLen)
+        all_data=[]
+        all_masks=[]
+        all_labels=[]
+        all_transforms=[]
+
+        for sentence in sentences:
+            tok_ids=[]
+            input_mask=[]
+            labels=[]
+            transform=[]
+
+            all_toks=[]
+            n=0
+            for idx, word in enumerate(sentence):
+                toks=self._tokenize(word)
+                all_toks.append(toks)
+                n+=len(toks)
+
+            cur=0
+            for idx, word in enumerate(sentence):
+                toks=all_toks[idx]
+                ind=list(np.zeros(n))
+                for j in range(cur,cur+len(toks)):
+                    ind[j]=1./len(toks)
+                cur+=len(toks)
+                transform.append(ind)
+
+                tok_ids.extend(self.convert_tokens_to_ids(toks))
+
+                input_mask.extend(np.ones(len(toks)))
+                labels.append(1)
+
+            all_data.append(tok_ids)
+            all_masks.append(input_mask)
+            all_labels.append(labels)
+            all_transforms.append(transform)
+
+        lengths = np.array([len(l) for l in all_data])
+
+        # Note sequence must be ordered from shortest to longest so current_batch will work
+        ordering = np.argsort(lengths)
+
+        ordered_data = [None for i in range(len(all_data))]
+        ordered_masks = [None for i in range(len(all_data))]
+        ordered_labels = [None for i in range(len(all_data))]
+        ordered_transforms = [None for i in range(len(all_data))]
+
+
+        for i, ind in enumerate(ordering):
+            ordered_data[i] = all_data[ind]
+            ordered_masks[i] = all_masks[ind]
+            ordered_labels[i] = all_labels[ind]
+            ordered_transforms[i] = all_transforms[ind]
+
+        batched_data=[]
+        batched_mask=[]
+        batched_labels=[]
+        batched_transforms=[]
+
+        i=0
+        current_batch=max_batch
+
+        while i < len(ordered_data):
+
+            batch_data=ordered_data[i:i+current_batch]
+            batch_mask=ordered_masks[i:i+current_batch]
+            batch_labels=ordered_labels[i:i+current_batch]
+            batch_transforms=ordered_transforms[i:i+current_batch]
+
+            max_len = max([len(sent) for sent in batch_data])
+            max_label = max([len(label) for label in batch_labels])
+
+            for j in range(len(batch_data)):
+
+                blen=len(batch_data[j])
+                blab=len(batch_labels[j])
+
+                for k in range(blen, max_len):
+                    batch_data[j].append(0)
+                    batch_mask[j].append(0)
+                    for z in range(len(batch_transforms[j])):
+                        batch_transforms[j][z].append(0)
+
+                for k in range(blab, max_label):
+                    batch_labels[j].append(-100)
+
+                for k in range(len(batch_transforms[j]), max_label):
+                    batch_transforms[j].append(np.zeros(max_len))
+
+            batched_data.append(batch_data)
+            batched_mask.append(batch_mask)
+            batched_labels.append(batch_labels)
+            batched_transforms.append(batch_transforms)
+
+            #bsize=torch.FloatTensor(batch_transforms).shape
+
+            i+=current_batch
+
+            # adjust batch size; sentences are ordered from shortest to longest so decrease as they get longer
+            if max_len > 100:
+                current_batch=12
+            if max_len > 200:
+                current_batch=6
+
+        #print(len(batch_data), len(batch_mask), len(batch_transforms))
+        return torch.LongTensor(batched_data).squeeze(), torch.FloatTensor(batched_mask).squeeze(), torch.FloatTensor(batched_transforms).squeeze()
+
+    '''
+    
+    def _tokenize(self, text):
+        if not isinstance(text, list):
+            text = [text]
+
+        outputs = []
+        for sentence in text:
+            tokens = sentence.split(" ")
+            wp_tokens = []
+            for token in tokens:
+                if token in ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]:
+                    wp_tokens.append(token)
+                else:
+                    wp_toks = self.encoder.encode(token)
+                    for wp in wp_toks:
+                        wp_tokens.append(self.reverseVocab[wp + 5])
+
+            outputs.append(SimpleNamespace(
+                tokens=wp_tokens,
+                ids=torch.Tensor(self.convert_tokens_to_ids(wp_tokens))
+            ))
+        return outputs
+    
+    '''
+
+    def _tokenize(self, text):
+        tokens = text.split(" ")
+        wp_tokens = []
+        for token in tokens:
+
+            if token in {"[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"}:
+                wp_tokens.append(token)
+            else:
+
+                wp_toks = self.encoder.encode(token)
+
+                for wp in wp_toks:
+                    wp_tokens.append(self.reverseVocab[wp + 5])
+        #print(wp_tokens)
+        return wp_tokens
+
+def main():
+    model = LatinBERT("../../latinBert/latin_bert/models/latin_bert", tokenizerPath="./tokenizer/latin.subword.encoder")
+
+    sents = ["arma virumque cano", "arma gravi numero violentaque bella parabam"]
+
+
+    output = model(sents)
+    print("end", output.shape)
+
+if __name__ == "__main__":
+    main()

embedders/labse.py

@@ -0,0 +1,39 @@
+import torch
+from transformers import BertModel, BertTokenizerFast
+import torch.nn.functional as F
+
+def similarity(embeddings_1, embeddings_2):
+    normalized_embeddings_1 = F.normalize(embeddings_1, p=2)
+    normalized_embeddings_2 = F.normalize(embeddings_2, p=2)
+    return torch.matmul(
+        normalized_embeddings_1, normalized_embeddings_2.transpose(0, 1)
+    )
+
+class LaBSE:
+    def __init__(self):
+        self.tokenizer = BertTokenizerFast.from_pretrained("setu4993/LaBSE")
+        self.model = BertModel.from_pretrained("setu4993/LaBSE")
+        #self.model = self.model.to('cuda')
+        self.model.eval()
+
+    @torch.no_grad()
+    def __call__(self, sentences):
+        if not isinstance(sentences, list):
+            sentences = [sentences]
+        tokens = self.tokenizer(sentences, return_tensors="pt", padding=True)
+        #print(tokens.input_ids.shape, tokens.token_type_ids.shape, tokens.attention_mask.shape)
+        #tokens = tokens.to("cuda")
+        outputs = self.model(**tokens)
+        embeddings = outputs.pooler_output
+        return F.normalize(embeddings, p=2).cpu()#.numpy()
+
+    @property
+    def dim(self):
+        return 768
+
+if __name__ == "__main__":
+    model = LaBSE()
+    sents = ["arma virumque cano", "arma gravi numero violentaque bella parabam"]
+
+    output = model(sents)
+    print("end", output.shape)

embedders/latin_bert/latin_bert/config.json

@@ -0,0 +1,19 @@
+{
+  "attention_probs_dropout_prob": 0.1,
+  "directionality": "bidi",
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "max_position_embeddings": 512,
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pooler_fc_size": 768,
+  "pooler_num_attention_heads": 12,
+  "pooler_num_fc_layers": 3,
+  "pooler_size_per_head": 128,
+  "pooler_type": "first_token_transform",
+  "type_vocab_size": 2,
+  "vocab_size": 32900
+}

embedders/latin_bert/latin_bert/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ad931a030dffd79a6cc13eaf5108352418285d3a04b510cc70b41fdd609635b5
+size 447786794