diff --git a/rag/nlp/synonym.py b/rag/nlp/synonym.py index ba9f0a04a..a826a3266 100644 --- a/rag/nlp/synonym.py +++ b/rag/nlp/synonym.py @@ -66,7 +66,7 @@ class Dealer: except Exception as e: logging.error("Fail to load synonym!" + str(e)) - def lookup(self, tk): + def lookup(self, tk, topn=8): if re.match(r"[a-z]+$", tk): res = list(set([re.sub("_", " ", syn.name().split(".")[0]) for syn in wordnet.synsets(tk)]) - set([tk])) return [t for t in res if t] @@ -76,7 +76,7 @@ class Dealer: res = self.dictionary.get(re.sub(r"[ \t]+", " ", tk.lower()), []) if isinstance(res, str): res = [res] - return res + return res[:topn] if __name__ == '__main__': diff --git a/rag/raptor.py b/rag/raptor.py index 6a6013605..eea7ae279 100644 --- a/rag/raptor.py +++ b/rag/raptor.py @@ -39,6 +39,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval: if response: return response response = self._llm_model.chat(system, history, gen_conf) + response = re.sub(r".*", "", response, flags=re.DOTALL) if response.find("**ERROR**") >= 0: raise Exception(response) set_llm_cache(self._llm_model.llm_name, system, response, history, gen_conf)