Update cross_encoder_reranking_train.py

cross_encoder_reranking_train.py

@@ -276,24 +276,29 @@ def hybrid_score(cross_encoder_score, semantic_score, weight_cross=0.7, weight_s
 def cross_encoder_reranking(query_text, doc_texts, model, tokenizer, batch_size=64, max_length=2048):
     device = next(model.parameters()).device
     cross_scores = []
-    query_emb = embed_text_list([query_text])[0]
+    query_emb = embed_text_list([query_text])[0]  # Move embedder to CPU
 
     instructed_query = get_detailed_instruct("", query_text)
 
-    for i in tqdm(range(0, len(doc_texts), batch_size), desc="Scoring documents", leave=False):
-        batch_docs = doc_texts[i:i+batch_size]
+    # Pre-create all input pairs (concatenation-based cross-encoder setup)
+    input_texts = [f"{instructed_query} {doc}" for doc in doc_texts]
 
-        input_texts = [instructed_query] + batch_docs
+    for i in tqdm(range(0, len(input_texts), batch_size), desc="Scoring documents", leave=False):
+        batch_input_texts = input_texts[i:i+batch_size]
 
         with torch.no_grad():
-            batch_dict = tokenizer(input_texts, max_length=max_length, padding=True, truncation=True, return_tensors='pt').to(device)
-
-            outputs = model(**batch_dict)
-            embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
-            embeddings = F.normalize(embeddings, p=2, dim=1)
-
-            batch_cross_scores = (embeddings[0].unsqueeze(0) @ embeddings[1:].T).squeeze(0).cpu().numpy()
-            cross_scores.extend(batch_cross_scores)
+            batch_dict = tokenizer(batch_input_texts, max_length=max_length, padding=True, truncation=True, return_tensors='pt').to(device)
+
+            # Mixed precision for faster inference and lower memory
+            with torch.cuda.amp.autocast():
+                outputs = model(**batch_dict)
+                embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
+                embeddings = F.normalize(embeddings, p=2, dim=1)
+
+                # Since queries are repeated in each pair, compare to instructed query embedding (first one)
+                query_vector = embeddings[0].unsqueeze(0)  # Use first as query
+                batch_cross_scores = (query_vector @ embeddings.T).squeeze(0).cpu().numpy()[1:]  # Exclude self-comparison
+                cross_scores.extend(batch_cross_scores)
 
     # Semantic scores
     doc_embeddings = embed_text_list(doc_texts)