Update README.md

README.md

@@ -100,6 +100,40 @@ To perform inference with the original model files, you’ll first need to insta
 After installation, download the model weights and use the Scaling inference module to load the 
 checkpoint, vocabulary, and configuration files.
 
+```
+from pathlib import Path
+from torch.nn import CosineSimilarity
+from scaling.transformer.inference import TransformerInferenceModule
+MODEL_PATH = "/path/to/model"
+inference_model = TransformerInferenceModule.from_checkpoint(
+    checkpoint_dir=Path(MODEL_PATH),
+)
+# embed the query:
+query = "Which country is Galileo from?"
+query_embeddings = inference_model.encode_queries(query, convert_to_tensor=True)
+print(f"Type of embeddings: {type(query_embeddings)},\n\
+       shape of query embeddings: {query_embeddings.shape}")
+# embed the documents:
+document_1 =  "Galileo is a German television program series produced and broadcast on ProSieben television network. It is also sold to broadcasters in other countries (namely Russia and Poland). The first show was broadcast in 1998, and is now stored in the Arctic World Archive in Svalbard, Norway, after being transferred to special film created by Piql."
+document_embeddings_1 = inference_model.encode_corpus(document_1, convert_to_tensor=True)
+document_2 = "Galileo di Vincenzo Bonaiuti de' Galilei (15 February 1564 - 8 January 1642), commonly referred to as Galileo Galilei or mononymously as Galileo, was an Italian (Florentine) astronomer, physicist and engineer, sometimes described as a polymath. He was born in the city of Pisa, then part of the Duchy of Florence and present-day Italy."
+document_embeddings_2 = inference_model.encode_corpus(document_2, convert_to_tensor=True)
+# customized embeddings steering the query:
+instruction = "Represent the question about TV shows to find a paragraph that answers it."
+steered_query_embeddings = inference_model.encode_queries(query, 
+                                                          instruction=instruction,
+                                                          convert_to_tensor=True)
+# compute similarity between steered query and both documents
+cossim = CosineSimilarity(dim=1, eps=1e-6)
+sim1 = round(cossim(document_embeddings_1, steered_query_embeddings).item(), 3)
+sim2 = round(cossim(document_embeddings_2, steered_query_embeddings).item(), 3)
+print("Steered embedding causes higher similarity of query to TV show:")
+print(f"Similarity query/TV show ({sim1}) > similarity query/Italian polymath: ({sim2})")
+```
+
+
+
+
 ## Bias, Risks, and Limitations
 
 <!-- This section is meant to convey both technical and sociotechnical limitations. -->