Update README.md

README.md

@@ -6,6 +6,9 @@ tags:
 - sentence-similarity
 datasets:
 - biu-nlp/abstract-sim
+widgets:
+- sentence-similarity
+- feature-extraction
 ---
 
 A model for mapping abstract sentence descriptions to sentences that fit the descriptions. Trained on Wikipedia. Use ```load_finetuned_model``` to load the query and sentence encoder, and ```encode_batch()``` to encode a sentence with the model.
@@ -14,6 +17,7 @@ A model for mapping abstract sentence descriptions to sentences that fit the des
 
 from transformers import AutoTokenizer, AutoModel
 import torch
+from typing import List
 
 def load_finetuned_model():
 
@@ -25,7 +29,7 @@ def load_finetuned_model():
         return tokenizer, query_encoder, sentence_encoder
 
 
-def encode_batch(model, tokenizer, sentences, device):
+def encode_batch(model, tokenizer, sentences: List[str], device: str):
     input_ids = tokenizer(sentences, padding=True, max_length=512, truncation=True, return_tensors="pt",
                           add_special_tokens=True).to(device)
     features = model(**input_ids)[0]