Merge pull request #8 from lukfre/decode-fix

Decode fix
IAAR-Shanghai · Dec 25, 2024 · c6ab86f · c6ab86f
2 parents 491160e + 6d78835
commit c6ab86f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/xfinder/modules/answer_comparator.py b/src/xfinder/modules/answer_comparator.py
@@ -40,6 +40,8 @@ def compare(
                 if type(standard_answer_range) == str:
                     standard_answer_range_list = ast.literal_eval(
                         standard_answer_range)
+                else:  
+                    standard_answer_range_list = standard_answer_range
                 for option in standard_answer_range_list:
                     if option[0] == correct and \
                             extracted.strip().rstrip(".").lower() == option[1].strip().rstrip(".").lower():

diff --git a/src/xfinder/modules/answer_extractor.py b/src/xfinder/modules/answer_extractor.py
@@ -126,8 +126,8 @@ def _execute_local_inference(self, query: Dict[str, Any]) -> str:
         output_ids = self.model.generate(
             input_ids, max_new_tokens=self.max_tokens, temperature=self.temperature)
         response = self.tokenizer.decode(
-            output_ids[0], skip_special_tokens=True)
-        return response.replace(prompt, '').strip()
+            output_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
+        return response.strip()
 
     def generate_output(self, question, llm_output, standard_answer_range) -> str:
         formatted_query = f'Question: """{question}"""\n\nOutput sentences: """{llm_output}"""\n\nAnswer range: {standard_answer_range}\n\nKey extracted answer: '