Upload folder using huggingface_hub

README.md

@@ -274,6 +274,16 @@ with torch.no_grad():
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+def last_logit_pool(logits: torch.Tensor,
+                    attention_mask: torch.Tensor) -> torch.Tensor:
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return logits[:, -1]
+    else:
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = logits.shape[0]
+        return torch.stack([logits[i, sequence_lengths[i]] for i in range(batch_size)], dim=0)
+
 def get_inputs(pairs, tokenizer, prompt=None, max_length=1024):
     if prompt is None:
         prompt = "Predict whether passage B contains an answer to query A."
@@ -285,6 +295,8 @@ def get_inputs(pairs, tokenizer, prompt=None, max_length=1024):
                            return_tensors=None,
                            add_special_tokens=False)['input_ids']
     inputs = []
+    query_lengths = []
+    prompt_lengths = []
     for query, passage in pairs:
         query_inputs = tokenizer(f'A: {query}',
                                  return_tensors=None,
@@ -309,25 +321,39 @@ def get_inputs(pairs, tokenizer, prompt=None, max_length=1024):
         item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
         item['attention_mask'] = [1] * len(item['input_ids'])
         inputs.append(item)
+        query_lengths.append(len([tokenizer.bos_token_id] + query_inputs['input_ids'] + sep_inputs))
+        prompt_lengths.append(len(sep_inputs + prompt_inputs))
+        
     return tokenizer.pad(
             inputs,
             padding=True,
             max_length=max_length + len(sep_inputs) + len(prompt_inputs),
             pad_to_multiple_of=8,
             return_tensors='pt',
-    )
+    ), query_lengths, prompt_lengths
 
 tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2.5-gemma2-lightweight', trust_remote_code=True)
+tokenizer.padding_side = 'right'
 model = AutoModelForCausalLM.from_pretrained('BAAI/bge-reranker-v2.5-gemma2-lightweight', trust_remote_code=True)
 model = model.to('cuda')
 model.eval()
 
 pairs = [['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']]
 with torch.no_grad():
-    inputs = get_inputs(pairs, tokenizer).to(model.device)
-    all_scores = model(**inputs, return_dict=True, cutoff_layers=[28], compress_ratio=2, compress_layer=[24, 40])
-    all_scores = [scores[:, -1].view(-1, ).float() for scores in all_scores[0]]
-    print(all_scores)
+    inputs, query_lengths, prompt_lengths = get_inputs(pairs, tokenizer)
+    inputs = inputs.to(model.device)
+    outputs = model(**inputs,
+                    return_dict=True,
+                    cutoff_layers=[28],
+                    compress_ratio=2,
+                    compress_layer=[24, 40],
+                    query_lengths=query_lengths,
+                    prompt_lengths=prompt_lengths)
+    scores = []
+    for i in range(len(outputs.logits)):
+        logits = last_logit_pool(outputs.logits[i], outputs.attention_masks[i])
+        scores.append(logits.cpu().float().tolist())
+    print(scores)
 ```
 
 ## Evaluation