first commit

.gitignore

@@ -0,0 +1,4 @@
+/env
+/__pycache__/
+
+.env

README.md

@@ -1,6 +1,6 @@
 ---
 title: Serica Semantic Search
-emoji: 🐨
+emoji: 🔍
 colorFrom: indigo
 colorTo: pink
 sdk: streamlit
@@ -9,5 +9,3 @@ app_file: app.py
 pinned: false
 license: agpl-3.0
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,120 @@
+import os
+import streamlit as st
+from elasticsearch import Elasticsearch
+
+from embedders.labse import LaBSE
+
+def search():
+    status_indicator.write(f"Loading model {model_name}...")
+    model = globals()[model_name]()
+    
+    status_indicator.write(f"Computing query embeddings...")
+    query_vector = model(query)[0, :].tolist()
+    
+    status_indicator.write(f"Performing query...")
+    target_field = f"{model_name}_features"
+    results = es.search(
+        index="sentences",
+        query={
+            "script_score": {
+                "query": {"match_all": {}},
+                "script": {
+                    "source": f"cosineSimilarity(params.query_vector, '{target_field}') + 1.0",
+                    "params": {"query_vector": query_vector}
+                }
+            }
+        },
+        size=limit
+    )
+
+    for result in results["hits"]["hits"]:
+        sentence = result['_source']['sentence']
+        score =  result['_score']
+        document = result['_source']['document']
+        number = result['_source']['number']
+
+        previous = es.search(
+            index="sentences",
+            query={
+                "bool": {
+                    "must": [{
+                        "term": {
+                            "document": document
+                            }
+                        },{
+                        "range": {
+                            "number": {
+                            "gte": number-3,
+                            "lt": number,
+                                }
+                            }
+                        }
+                    ]
+                }
+            }
+        )
+
+        previous_hits = sorted(previous["hits"]["hits"], key=lambda e: e["_source"]["number"])
+        previous_context = "".join([r["_source"]["sentence"] for r in previous_hits])
+
+
+        subsequent = es.search(
+            index="sentences",
+            query={
+                "bool": {
+                    "must": [{
+                        "term": {
+                            "document": document
+                            }
+                        },{
+                        "range": {
+                            "number": {
+                            "lte": number+3,
+                            "gt": number,
+                                }
+                            }
+                        }
+                    ]
+                }
+            }
+        )
+
+        subsequent_hits = sorted(subsequent["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']}"
+
+        results_placeholder.markdown(f"#### {document_name} (score: {score:.2f})\n{previous_context} **{sentence}** {subsequent_context}")
+
+
+    status_indicator.write(f"Results ready...")
+
+es = Elasticsearch(os.environ["ELASTIC_HOST"], basic_auth=os.environ["ELASTIC_AUTH"].split(":"))
+
+st.header("Serica Semantic Search")
+st.write("Perform a semantic search using a Sentence Embedding Transformer model on the SERICA database")
+model_name = st.selectbox("Model", ["LaBSE"])
+limit = st.number_input("Number of results", 10)
+query = st.text_input("Query", value="")
+status_indicator = st.empty()
+do_search = st.button("Search")
+results_placeholder = st.container()
+
+if do_search:
+    search()

embedders/labse.py

@@ -0,0 +1,26 @@
+import torch
+from transformers import BertModel, BertTokenizerFast
+import torch.nn.functional as F
+
+class LaBSE:
+    def __init__(self):
+        self.tokenizer = BertTokenizerFast.from_pretrained("setu4993/LaBSE")
+        self.model = BertModel.from_pretrained("setu4993/LaBSE")
+        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)
+        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__":
+    labse = LaBSE()
+    print(labse(["odi et amo", "quare id faciam"]).shape)

requirements.txt

@@ -0,0 +1,19 @@
+certifi==2022.6.15
+charset-normalizer==2.1.0
+elastic-transport==8.1.2
+elasticsearch==8.3.3
+filelock==3.7.1
+huggingface-hub==0.8.1
+idna==3.3
+numpy==1.23.1
+packaging==21.3
+pyparsing==3.0.9
+PyYAML==6.0
+regex==2022.7.25
+requests==2.28.1
+tokenizers==0.12.1
+tqdm==4.64.0
+transformers==4.21.0
+typing-extensions==4.3.0
+urllib3==1.26.11
+torch==1.12.0