flashrag-dev 0.1.4.dev20250531__tar.gz → 0.1.4.dev20250610__tar.gz

{flashrag_dev-0.1.4.dev20250531 → flashrag_dev-0.1.4.dev20250610}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag_dev
-Version: 0.1.4.dev20250531
+Version: 0.1.4.dev20250610
 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.dev20250531 → flashrag_dev-0.1.4.dev20250610}/flashrag/prompt/base_prompt.py

@@ -23,16 +23,9 @@ class PromptTemplate:
             self.is_chat = False
             if "chat" in model_name or "instruct" in model_name:
                 self.is_chat = True
-            self.tokenizer = AutoTokenizer.from_pretrained(self.generator_path, trust_remote_code=True)
         else:
             self.is_chat = True
             self.enable_chat = True
-            try:
-                self.tokenizer = tiktoken.encoding_for_model(config['generator_model'])
-            except Exception as e:
-                print("Error: ", e)
-                warnings.warn("This model is not supported by tiktoken. Use gpt-3.5-turbo instead.")
-                self.tokenizer = tiktoken.encoding_for_model('gpt-3.5-turbo')
 
         if len(system_prompt) == 0 and len(user_prompt) == 0:
             system_prompt = self.base_system_prompt
@@ -41,8 +34,22 @@ class PromptTemplate:
         self.user_prompt = user_prompt
         self.enable_chat = enable_chat
         self.reference_template = reference_template
+        self.tokenizer = None
 
         # self._check_placeholder()
+    def _get_tokenizer(self):
+        if self.tokenizer is None:
+            if self.is_openai:
+                try:
+                    self.tokenizer = tiktoken.encoding_for_model(self.config['generator_model'])
+                except Exception as e:
+                    print("Error: ", e)
+                    warnings.warn("This model is not supported by tiktoken. Use gpt-3.5-turbo instead.")
+                    self.tokenizer = tiktoken.encoding_for_model('gpt-3.5-turbo')
+            else:
+                self.tokenizer = AutoTokenizer.from_pretrained(self.generator_path, trust_remote_code=True)
+        return self.tokenizer
+
 
     def _check_placeholder(self):
         # check placeholder in prompt
@@ -64,7 +71,7 @@ class PromptTemplate:
                 assert isinstance(prompt, list)
                 for message in prompt:
                     role_content = message['content']
-                    encoded_message = self.tokenizer.encode(role_content)
+                    encoded_message = self._get_tokenizer().encode(role_content)
 
                     if total_tokens + len(encoded_message) <= self.max_input_len:
                         truncated_messages.append(message)
@@ -72,29 +79,30 @@ class PromptTemplate:
                     else:
                         print(f"The input text length is greater than the maximum length ({total_tokens + len(encoded_message)} > {self.max_input_len}) and has been truncated!")
                         remaining_tokens = self.max_input_len - total_tokens
-                        truncated_message = self.encoding.decode(encoded_message[:remaining_tokens])
+                        truncated_message = self._get_tokenizer().decode(encoded_message[:remaining_tokens])
                         message['content'] = truncated_message
                         truncated_messages.append(message)
                         break
             else:
                 assert isinstance(prompt, str)
-                tokenized_prompt = self.tokenizer.encode(prompt,allowed_special={'<|endoftext|>'})
+                tokenized_prompt = self._get_tokenizer().encode(prompt,allowed_special={'<|endoftext|>'})
                 half = int(self.max_input_len / 2)
-                truncated_messages = self.tokenizer.decode(tokenized_prompt[:half]) + self.tokenizer.decode(tokenized_prompt[-half:])
+                truncated_messages = self._get_tokenizer().decode(tokenized_prompt[:half]) + self._get_tokenizer().decode(tokenized_prompt[-half:])
 
             return truncated_messages
 
         else:
-            if self.tokenizer is None:
-                self.tokenizer = AutoTokenizer.from_pretrained(self.generator_path, trust_remote_code=True)
             assert isinstance(prompt, str)
-            tokenized_prompt = self.tokenizer(prompt, truncation=False, return_tensors="pt").input_ids[0]
+            try:
+                tokenized_prompt = self._get_tokenizer().encode(prompt, truncation=False, return_tensors="pt").input_ids[0]
+            except:
+                tokenized_prompt = self._get_tokenizer().encode(prompt, truncation=False, return_tensors="pt")[0]
 
             if len(tokenized_prompt) > self.max_input_len:
                 print(f"The input text length is greater than the maximum length ({len(tokenized_prompt)} > {self.max_input_len}) and has been truncated!")
-                half = int(self.max_input_len / 2)
-                prompt = self.tokenizer.decode(tokenized_prompt[:half], skip_special_tokens=True) + \
-                        self.tokenizer.decode(tokenized_prompt[-half:], skip_special_tokens=True)
+                half = int(self.max_input_len / 2) - 20
+                prompt = self._get_tokenizer().decode(tokenized_prompt[:half], skip_special_tokens=True) + \
+                        self._get_tokenizer().decode(tokenized_prompt[-half:], skip_special_tokens=True)
             return prompt
 
 
@@ -107,7 +115,7 @@ class PromptTemplate:
                 if self.is_openai:
                     self.truncate_prompt(messages)
                 else:
-                    prompt = self.tokenizer.apply_chat_template(
+                    prompt = self._get_tokenizer().apply_chat_template(
                         messages, tokenize=False, add_generation_prompt=True
                     )
                     return self.truncate_prompt(prompt)
@@ -136,7 +144,11 @@ class PromptTemplate:
             if user_prompt != "":
                 input.append({"role": "user", "content": user_prompt})
             if not self.is_openai:
-                input = self.tokenizer.apply_chat_template(input, tokenize=False, add_generation_prompt=True)
+                try:
+                    input = self._get_tokenizer().apply_chat_template(input, tokenize=False, add_generation_prompt=True)
+                except:
+                    print("Warning: the generator tokenizer not support `apply_chat_template`")
+                    input = system_prompt + '\n\n' + user_prompt
         else:
             input = "\n\n".join([prompt for prompt in [system_prompt, user_prompt] if prompt != ""])
 
@@ -180,7 +192,7 @@ class PromptTemplate:
                     **params,
                 )
 
-                possible_prompt_tokens = tokenizer.encode(possible_prompt)
+                possible_prompt_tokens = self._get_tokenizer().encode(possible_prompt)
                 if len(possible_prompt_tokens) <= max_length:
                     final_examplars = examplars[:num]
                     break

flashrag_dev-0.1.4.dev20250610/flashrag/version.py

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

{flashrag_dev-0.1.4.dev20250531 → flashrag_dev-0.1.4.dev20250610}/flashrag_dev.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag-dev
-Version: 0.1.4.dev20250531
+Version: 0.1.4.dev20250610
 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.dev20250531/flashrag/version.py

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