flashrag-dev 0.1.4.dev20250423__tar.gz → 0.1.4.dev20250425__tar.gz

{flashrag_dev-0.1.4.dev20250423 → flashrag_dev-0.1.4.dev20250425}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag_dev
-Version: 0.1.4.dev20250423
+Version: 0.1.4.dev20250425
 Summary: A library for efficient Retrieval-Augmented Generation research
 Home-page: https://github.com/RUC-NLPIR/FlashRAG
 Author: Jiajie Jin, Yutao Zhu, Chenghao Zhang, Xinyu Yang, Zhicheng Dou

{flashrag_dev-0.1.4.dev20250423 → flashrag_dev-0.1.4.dev20250425}/flashrag/retriever/encoder.py

@@ -24,7 +24,7 @@ class Encoder:
             Encodes a list of queries into embeddings.
     """
 
-    def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16, instruction, silent=False):
+    def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16=True, instruction="", silent=False):
         self.model_name = model_name
         self.model_path = model_path
         self.pooling_method = pooling_method
@@ -44,7 +44,9 @@ class Encoder:
         )
         inputs = {k: v.cuda() for k, v in inputs.items()}
 
-        if "T5" in type(self.model).__name__ or (isinstance(self.model, torch.nn.DataParallel) and "T5" in type(self.model.module).__name__):
+        if "T5" in type(self.model).__name__ or (
+            isinstance(self.model, torch.nn.DataParallel) and "T5" in type(self.model.module).__name__
+        ):
             # T5-based retrieval model
             decoder_input_ids = torch.zeros((inputs["input_ids"].shape[0], 1), dtype=torch.long).to(
                 inputs["input_ids"].device
@@ -54,11 +56,9 @@ class Encoder:
 
         else:
             output = self.model(**inputs, return_dict=True)
-            pooler_output = output.get('pooler_output', None)
-            last_hidden_state = output.get('last_hidden_state', None)
-            query_emb = pooling(
-                pooler_output, last_hidden_state, inputs["attention_mask"], self.pooling_method
-            )
+            pooler_output = output.get("pooler_output", None)
+            last_hidden_state = output.get("last_hidden_state", None)
+            query_emb = pooling(pooler_output, last_hidden_state, inputs["attention_mask"], self.pooling_method)
         if "dpr" not in self.model_name:
             query_emb = torch.nn.functional.normalize(query_emb, dim=-1)
         query_emb = query_emb.detach().cpu().numpy()
@@ -117,7 +117,11 @@ class STEncoder:
     def encode(self, query_list: Union[List[str], str], batch_size=64, is_query=True) -> np.ndarray:
         query_list = parse_query(self.model_name, query_list, self.instruction, is_query)
         query_emb = self.model.encode(
-            query_list, batch_size=batch_size, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=not self.silent
+            query_list,
+            batch_size=batch_size,
+            convert_to_numpy=True,
+            normalize_embeddings=True,
+            show_progress_bar=not self.silent,
         )
         query_emb = query_emb.astype(np.float32, order="C")
 
@@ -173,11 +177,10 @@ class ClipEncoder:
         # set model max length for model that not specified in config.json
         if self.processor is not None and self.processor.tokenizer.model_max_length > 100000:
             try:
-                model_max_length = config['text_config']['max_position_embeddings']
+                model_max_length = config["text_config"]["max_position_embeddings"]
             except:
                 model_max_length = 512
-            self.processor.tokenizer.model_max_length = model_max_length    
-
+            self.processor.tokenizer.model_max_length = model_max_length
 
     @torch.inference_mode()
     def single_batch_encode(self, query_list: Union[List[str], str], modal="image") -> np.ndarray:
@@ -207,7 +210,7 @@ class ClipEncoder:
     @torch.inference_mode()
     def encode_image(self, image_list: List) -> np.ndarray:
         # Each item in image_list: PIL Image, local path, or URL
-        if self.model_type == "CLIPModel" or self.model_type == 'ChineseCLIPModel':
+        if self.model_type == "CLIPModel" or self.model_type == "ChineseCLIPModel":
             # need handle image
             image_list = [parse_image(image) for image in image_list]
             inputs = self.processor(images=image_list, return_tensors="pt")
@@ -224,7 +227,7 @@ class ClipEncoder:
     @torch.inference_mode()
     def encode_text(self, text_list: List[str]) -> np.ndarray:
         # Each item in image_list: PIL Image, local path, or URL
-        if self.model_type == "CLIPModel" or self.model_type == 'ChineseCLIPModel':
+        if self.model_type == "CLIPModel" or self.model_type == "ChineseCLIPModel":
             inputs = self.processor(
                 text=text_list,
                 padding=True,

{flashrag_dev-0.1.4.dev20250423 → flashrag_dev-0.1.4.dev20250425}/flashrag/retriever/retriever.py

@@ -305,7 +305,8 @@ class BM25Retriever(BaseTextRetriever):
             else:
                 results = load_docs(self.corpus, [hit.docid for hit in hits])
         elif self.backend == "bm25s":
-            query_tokens = self.tokenizer.tokenize([query], return_as="tuple", update_vocab=False)
+            # query_tokens = self.tokenizer.tokenize([query], return_as="tuple", update_vocab=False)
+            query_tokens = bm25s.tokenize([query])
             results, scores = self.searcher.retrieve(query_tokens, k=num)
             results = list(results[0])
             scores = list(scores[0])
@@ -327,7 +328,8 @@ class BM25Retriever(BaseTextRetriever):
                 results.append(item_result)
                 scores.append(item_score)
         elif self.backend == "bm25s":
-            query_tokens = self.tokenizer.tokenize(query, return_as="tuple", update_vocab=False)
+            # query_tokens = self.tokenizer.tokenize(query, return_as="tuple", update_vocab=False)
+            query_tokens = bm25s.tokenize(query)
             results, scores = self.searcher.retrieve(query_tokens, k=num)
         else:
             assert False, "Invalid bm25 backend!"

flashrag_dev-0.1.4.dev20250425/flashrag/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.4dev20250425"

{flashrag_dev-0.1.4.dev20250423 → flashrag_dev-0.1.4.dev20250425}/flashrag_dev.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag-dev
-Version: 0.1.4.dev20250423
+Version: 0.1.4.dev20250425
 Summary: A library for efficient Retrieval-Augmented Generation research
 Home-page: https://github.com/RUC-NLPIR/FlashRAG
 Author: Jiajie Jin, Yutao Zhu, Chenghao Zhang, Xinyu Yang, Zhicheng Dou

flashrag_dev-0.1.4.dev20250423/flashrag/version.py

@@ -1 +0,0 @@
-__version__ = "0.1.4dev20250423"