Update README.md

README.md

@@ -14,40 +14,46 @@ Below is an example to encode a query and a passage, and then compute their simi
 import torch
 from transformers import AutoModel, AutoTokenizer, LlamaModel
 
-def get_query_inputs(query, tokenizer, max_length=512):
+def get_query_inputs(queries, tokenizer, max_length=512):
     prefix = '"'
     suffix = '", predict the following passage within eight words: <s9><s10><s11><s12><s13><s14><s15><s16>'
     prefix_ids = tokenizer(prefix, return_tensors=None)['input_ids']
     suffix_ids = tokenizer(suffix, return_tensors=None)['input_ids'][1:]
-    inputs = tokenizer(query,
-                       return_tensors=None,
-                       max_length=max_length,
-                       truncation=True,
-                       add_special_tokens=False)
-    inputs['input_ids'] = prefix_ids + inputs['input_ids'] + suffix_ids
-    inputs['attention_mask'] = [1] * len(inputs['input_ids'])
+    queries_inputs = []
+    for query in queries:
+        inputs = tokenizer(query,
+                           return_tensors=None,
+                           max_length=max_length,
+                           truncation=True,
+                           add_special_tokens=False)
+        inputs['input_ids'] = prefix_ids + inputs['input_ids'] + suffix_ids
+        inputs['attention_mask'] = [1] * len(inputs['input_ids'])
+        queries_inputs.append(inputs)
     return tokenizer.pad(
-            [inputs],
+            queries_inputs,
             padding=True,
             max_length=max_length,
             pad_to_multiple_of=8,
             return_tensors='pt',
         )
 
-def get_passage_inputs(passage, tokenizer, max_length=512):
+def get_passage_inputs(passages, tokenizer, max_length=512):
     prefix = '"'
     suffix = '", summarize the above passage within eight words: <s1><s2><s3><s4><s5><s6><s7><s8>'
     prefix_ids = tokenizer(prefix, return_tensors=None)['input_ids']
     suffix_ids = tokenizer(suffix, return_tensors=None)['input_ids'][1:]
-    inputs = tokenizer(passage,
-                       return_tensors=None,
-                       max_length=max_length,
-                       truncation=True,
-                       add_special_tokens=False)
-    inputs['input_ids'] = prefix_ids + inputs['input_ids'] + suffix_ids
-    inputs['attention_mask'] = [1] * len(inputs['input_ids'])
+    passages_inputs = []
+    for passage in passages:
+        inputs = tokenizer(passage,
+                           return_tensors=None,
+                           max_length=max_length,
+                           truncation=True,
+                           add_special_tokens=False)
+        inputs['input_ids'] = prefix_ids + inputs['input_ids'] + suffix_ids
+        inputs['attention_mask'] = [1] * len(inputs['input_ids'])
+        passages_inputs.append(inputs)
     return tokenizer.pad(
-            [inputs],
+            passages_inputs,
             padding=True,
             max_length=max_length,
             pad_to_multiple_of=8,
@@ -62,8 +68,8 @@ model = AutoModel.from_pretrained('cfli/LLARA-beir')
 query = "What is llama?"
 title = "Llama"
 passage = "The llama is a domesticated South American camelid, widely used as a meat and pack animal by Andean cultures since the pre-Columbian era."
-query_input = get_query_inputs(query, tokenizer)
-passage_input = get_passage_inputs(passage, tokenizer)
+query_input = get_query_inputs([query], tokenizer)
+passage_input = get_passage_inputs([passage], tokenizer)
 
 
 with torch.no_grad():
@@ -84,4 +90,5 @@ with torch.no_grad():
     print(score)
 
 
+
 ```