Update README.md

README.md

@@ -37,12 +37,13 @@ def encode_batch(model, tokenizer, sentences: List[str], device: str):
     features =  torch.sum(features[:,1:,:] * input_ids["attention_mask"][:,1:].unsqueeze(-1), dim=1) / torch.clamp(torch.sum(input_ids["attention_mask"][:,1:], dim=1, keepdims=True), min=1e-9)
     return features
 
+```
 
+Usage example:
 
-if __name__ == "__main__":
-
-    tokenizer, query_encoder, sentence_encoder = load_finetuned_model()
-    relevant_sentences = ["Fingersoft's parent company is the Finger Group.",
+```python
+tokenizer, query_encoder, sentence_encoder = load_finetuned_model()
+relevant_sentences = ["Fingersoft's parent company is the Finger Group.",
                       "WHIRC – a subsidiary company of Wright-Hennepin",
                       "CK Life Sciences International (Holdings) Inc. (), or CK Life Sciences, is a subsidiary of CK Hutchison Holdings",
                       "EM Microelectronic-Marin (subsidiary of The Swatch Group).",
@@ -50,7 +51,7 @@ if __name__ == "__main__":
                       "Volt Technical Resources is a business unit of Volt Workforce Solutions, a subsidiary of Volt Information Sciences (currently trading over-the-counter as VISI.)."
              ]
 
-    irrelevant_sentences = ["The second company is deemed to be a subsidiary of the parent company.",
+irrelevant_sentences = ["The second company is deemed to be a subsidiary of the parent company.",
                         "The company has gone through more than one incarnation.",
                         "The company is owned by its employees.",
                         "Larger companies compete for market share by acquiring smaller companies that may own a particular market sector.",
@@ -58,19 +59,34 @@ if __name__ == "__main__":
                         "It is a holding company that provides services through its subsidiaries in the following areas: oil and gas, industrial and infrastructure, government and power."
                         ]
 
-    all_sentences = relevant_sentences + irrelevant_sentences
-    query = "<query>: A company is a part of a larger company."
+all_sentences = relevant_sentences + irrelevant_sentences
+query = "<query>: A company is a part of a larger company."
     
-    embeddings = encode_batch(sentence_encoder, tokenizer, all_sentences, "cpu").detach().cpu().numpy()
-    query_embedding = encode_batch(query_encoder, tokenizer, [query], "cpu").detach().cpu().numpy()
-
-    sims = cosine_similarity(query_embedding, embeddings)[0]
-    sentences_sims = list(zip(all_sentences, sims))
-    sentences_sims.sort(key=lambda x: x[1], reverse=True)
-
-    for s, sim in sentences_sims:
-        print(s, sim)
-
-  
-
+embeddings = encode_batch(sentence_encoder, tokenizer, all_sentences, "cpu").detach().cpu().numpy()
+query_embedding = encode_batch(query_encoder, tokenizer, [query], "cpu").detach().cpu().numpy()
+
+sims = cosine_similarity(query_embedding, embeddings)[0]
+sentences_sims = list(zip(all_sentences, sims))
+sentences_sims.sort(key=lambda x: x[1], reverse=True)
+
+for s, sim in sentences_sims:
+    print(s, sim)
+
+```
+
+Expected output:
+
+```
+WHIRC – a subsidiary company of Wright-Hennepin 0.9396286
+EM Microelectronic-Marin (subsidiary of The Swatch Group). 0.93929046
+Fingersoft's parent company is the Finger Group. 0.936247
+CK Life Sciences International (Holdings) Inc. (), or CK Life Sciences, is a subsidiary of CK Hutchison Holdings 0.9350312
+The company is currently a division of the corporate group Jam Industries. 0.9273489
+Volt Technical Resources is a business unit of Volt Workforce Solutions, a subsidiary of Volt Information Sciences (currently trading over-the-counter as VISI.). 0.9005086
+The second company is deemed to be a subsidiary of the parent company. 0.6723645
+It is a holding company that provides services through its subsidiaries in the following areas: oil and gas, industrial and infrastructure, government and power. 0.60081375
+A parent company is a company that owns 51% or more voting stock in another firm (or subsidiary). 0.59490484
+The company is owned by its employees. 0.55286574
+The company has gone through more than one incarnation. 0.38889483
+Larger companies compete for market share by acquiring smaller companies that may own a particular market sector. 0.25472647
 ```