feat(llm): BREAKING CHANGE update keyword extraction method

hugegraph-llm/.gitignore

@@ -1,7 +1,7 @@
 src/hugegraph_llm/resources/*
-!/src/hugegraph_llm/resources/demo/*
-!/src/hugegraph_llm/resources/nltk_data/*
-!/src/hugegraph_llm/resources/prompt_examples/*
+!/src/hugegraph_llm/resources/demo/
+!/src/hugegraph_llm/resources/nltk_data/corpora/stopwords/
+!/src/hugegraph_llm/resources/prompt_examples/
 
 
 uv.lock

hugegraph-llm/pyproject.toml

@@ -39,6 +39,9 @@ dependencies = [
     "numpy",
     "pandas",
     "pydantic",
+    "scipy",
+    "python-igraph",
+
 
     # LLM specific dependencies
     "openai",

hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py

@@ -46,6 +46,7 @@ class RAGRequest(BaseModel):
     topk_return_results: int = Query(20, description="Number of sorted results to return finally.")
     vector_dis_threshold: float = Query(0.9, description="Threshold for vector similarity\
                                          (results greater than this will be ignored).")
+    extract_method: str = Query("Hybrid", description="Method to extract keywords from the text.")
     topk_per_keyword: int = Query(1, description="TopK results returned for each keyword \
                                    extracted from the query, by default only the most similar one is returned.")
     client_config: Optional[GraphConfigRequest] = Query(None, description="hugegraph server config.")
@@ -56,6 +57,10 @@ class RAGRequest(BaseModel):
         prompt.keywords_extract_prompt,
         description="Prompt for extracting keywords from query.",
     )
+    mask_words: Optional[str] = Query(
+        prompt.maskword_input_text,
+        description="Mask words to be protected during word segmentation.",
+    )
     gremlin_tmpl_num: int = Query(1, description="Number of Gremlin templates to use.")
     gremlin_prompt: Optional[str] = Query(
         prompt.gremlin_generate_prompt,

hugegraph-llm/src/hugegraph_llm/api/rag_api.py

@@ -27,11 +27,12 @@
     RerankerConfigRequest,
     GraphRAGRequest,
 )
-from hugegraph_llm.config import huge_settings
 from hugegraph_llm.api.models.rag_response import RAGResponse
+from hugegraph_llm.config import huge_settings
 from hugegraph_llm.config import llm_settings, prompt
 from hugegraph_llm.utils.log import log
 
+
 # pylint: disable=too-many-statements
 def rag_http_api(
     router: APIRouter,
@@ -60,9 +61,11 @@ def rag_answer_api(req: RAGRequest):
             topk_return_results=req.topk_return_results,
             vector_dis_threshold=req.vector_dis_threshold,
             topk_per_keyword=req.topk_per_keyword,
+            keywords_extract_method=req.extract_method,
             # Keep prompt params in the end
             custom_related_information=req.custom_priority_info,
             answer_prompt=req.answer_prompt or prompt.answer_prompt,
+            mask_words=req.mask_words or prompt.maskword_input_text,
             keywords_extract_prompt=req.keywords_extract_prompt or prompt.keywords_extract_prompt,
             gremlin_prompt=req.gremlin_prompt or prompt.gremlin_generate_prompt,
         )

hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py

@@ -15,8 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import sys
 import os
+import sys
 from pathlib import Path
 
 import yaml
@@ -49,6 +49,7 @@ class BasePromptConfig:
     doc_input_text: str = ""
     _language_generated: str = ""
     generate_extract_prompt_template: str = ""
+    maskword_input_text: str = ""
 
     def ensure_yaml_file_exists(self):
         current_dir = Path.cwd().resolve()
@@ -117,6 +118,7 @@ def to_literal(val):
             "doc_input_text": to_literal(self.doc_input_text),
             "_language_generated": str(self.llm_settings.language).lower().strip(),
             "generate_extract_prompt_template": to_literal(self.generate_extract_prompt_template),
+            "maskword_input_text": to_literal(self.maskword_input_text),
         }
         with open(yaml_file_path, "w", encoding="utf-8") as file:
             yaml.dump(data, file, allow_unicode=True, sort_keys=False, default_flow_style=False)

hugegraph-llm/src/hugegraph_llm/config/prompt_config.py

@@ -427,4 +427,10 @@ def __init__(self, llm_config_object):
 ## Your Generated "Graph Extract Prompt Header":
 ## Language Requirement:
 Please generate the prompt in {language} language.
+"""
+
+    maskword_input_text: str = r"""/'https?://\S+|www\.\S+/,
+            /\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b/,
+            /\b\w+(?:[-’\']\w+)+\b/,
+            /\b\d+[,.]\d+\b/
 """

hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py

@@ -16,6 +16,7 @@
 # under the License.
 
 import argparse
+
 import gradio as gr
 import uvicorn
 from fastapi import FastAPI, Depends, APIRouter
@@ -101,6 +102,7 @@ def init_rag_ui() -> gr.Interface:
                 textbox_answer_prompt_input,
                 textbox_keywords_extract_prompt_input,
                 textbox_custom_related_information,
+                textbox_textrank_mask_words,
             ) = create_rag_block()
         with gr.Tab(label="3. Text2gremlin ⚙️"):
             textbox_gremlin_inp, textbox_gremlin_schema, textbox_gremlin_prompt = (
@@ -129,6 +131,7 @@ def refresh_ui_config_prompt() -> tuple:
                 prompt.answer_prompt,
                 prompt.keywords_extract_prompt,
                 prompt.custom_rerank_info,
+                prompt.maskword_input_text,
                 prompt.default_question,
                 huge_settings.graph_name,
                 prompt.gremlin_generate_prompt,
@@ -149,6 +152,7 @@ def refresh_ui_config_prompt() -> tuple:
                 textbox_answer_prompt_input,
                 textbox_keywords_extract_prompt_input,
                 textbox_custom_related_information,
+                textbox_textrank_mask_words,
                 textbox_gremlin_inp,
                 textbox_gremlin_schema,
                 textbox_gremlin_prompt,

hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py

@@ -18,18 +18,19 @@
 # pylint: disable=E1101
 
 import os
-from typing import AsyncGenerator, Tuple, Literal, Optional
+from typing import AsyncGenerator, Literal, Optional, Tuple
 
 import gradio as gr
 import pandas as pd
 from gradio.utils import NamedString
 
-from hugegraph_llm.config import resource_path, prompt, huge_settings, llm_settings
+from hugegraph_llm.config import huge_settings, llm_settings, prompt, resource_path
 from hugegraph_llm.operators.graph_rag_task import RAGPipeline
-from hugegraph_llm.utils.decorators import with_task_id
 from hugegraph_llm.operators.llm_op.answer_synthesize import AnswerSynthesize
+from hugegraph_llm.utils.decorators import with_task_id
 from hugegraph_llm.utils.log import log
 
+
 def rag_answer(
     text: str,
     raw_answer: bool,
@@ -42,6 +43,9 @@ def rag_answer(
     custom_related_information: str,
     answer_prompt: str,
     keywords_extract_prompt: str,
+    keywords_extract_method: str,
+    mask_words: str,
+    max_keywords_num: int = 5,
     gremlin_tmpl_num: Optional[int] = -1,
     gremlin_prompt: Optional[str] = None,
     max_graph_items=30,
@@ -66,6 +70,7 @@ def rag_answer(
         keywords_extract_prompt,
         text,
         vector_only_answer,
+        mask_words
     )
     if raw_answer is False and not vector_search and not graph_search:
         gr.Warning("Please select at least one generate mode.")
@@ -75,7 +80,12 @@ def rag_answer(
     if vector_search:
         rag.query_vector_index()
     if graph_search:
-        rag.extract_keywords(extract_template=keywords_extract_prompt).keywords_to_vid(
+        rag.extract_keywords(
+            extract_template=keywords_extract_prompt,
+            max_keywords=max_keywords_num,
+            extract_method=keywords_extract_method,
+            mask_words=mask_words
+        ).keywords_to_vid(
             vector_dis_threshold=vector_dis_threshold,
             topk_per_keyword=topk_per_keyword,
         ).import_schema(huge_settings.graph_name).query_graphdb(
@@ -125,6 +135,7 @@ def update_ui_configs(
     keywords_extract_prompt,
     text,
     vector_only_answer,
+    mask_words,
 ):
     gremlin_prompt = gremlin_prompt or prompt.gremlin_generate_prompt
     should_update_prompt = (
@@ -133,13 +144,15 @@ def update_ui_configs(
         or prompt.keywords_extract_prompt != keywords_extract_prompt
         or prompt.gremlin_generate_prompt != gremlin_prompt
         or prompt.custom_rerank_info != custom_related_information
+        or prompt.maskword_input_text != mask_words
     )
     if should_update_prompt:
         prompt.custom_rerank_info = custom_related_information
         prompt.default_question = text
         prompt.answer_prompt = answer_prompt
         prompt.keywords_extract_prompt = keywords_extract_prompt
         prompt.gremlin_generate_prompt = gremlin_prompt
+        prompt.maskword_input_text = mask_words
         prompt.update_yaml_file()
     vector_search = vector_only_answer or graph_vector_answer
     graph_search = graph_only_answer or graph_vector_answer
@@ -157,6 +170,9 @@ async def rag_answer_streaming(
     custom_related_information: str,
     answer_prompt: str,
     keywords_extract_prompt: str,
+    keywords_extract_method: str,
+    mask_words: str,
+    max_keywords_num: int = 5,
     gremlin_tmpl_num: Optional[int] = -1,
     gremlin_prompt: Optional[str] = None,
 ) -> AsyncGenerator[Tuple[str, str, str, str], None]:
@@ -177,6 +193,7 @@ async def rag_answer_streaming(
         keywords_extract_prompt,
         text,
         vector_only_answer,
+        mask_words
     )
     if raw_answer is False and not vector_search and not graph_search:
         gr.Warning("Please select at least one generate mode.")
@@ -187,7 +204,12 @@ async def rag_answer_streaming(
     if vector_search:
         rag.query_vector_index()
     if graph_search:
-        rag.extract_keywords(extract_template=keywords_extract_prompt).keywords_to_vid().import_schema(
+        rag.extract_keywords(
+            extract_template=keywords_extract_prompt,
+            extract_method=keywords_extract_method,
+            mask_words=mask_words,
+            max_keywords=max_keywords_num
+        ).keywords_to_vid().import_schema(
             huge_settings.graph_name
         ).query_graphdb(
             num_gremlin_generate_example=gremlin_tmpl_num,
@@ -261,7 +283,6 @@ def create_rag_block():
                 show_copy_button=True,
                 latex_delimiters=[{"left": "$", "right": "$", "display": False}],
             )
-
             answer_prompt_input = gr.Textbox(
                 value=prompt.answer_prompt, label="Query Prompt", show_copy_button=True, lines=7
             )
@@ -271,6 +292,15 @@ def create_rag_block():
                 show_copy_button=True,
                 lines=7,
             )
+            mask_words_input = gr.Textbox(
+                value=prompt.maskword_input_text,
+                label="TextRank mask words",
+                info=r"""Enter words or regular expressions to protect during word segmentation.
+                    Separate items with a comma and enclose regular expressions in "/".
+                    (Example: C++,/https?://\S+|www\.\S+/,//)""",
+                show_copy_button=True,
+                lines=7,
+            )
 
         with gr.Column(scale=1):
             with gr.Row():
@@ -279,6 +309,11 @@ def create_rag_block():
             with gr.Row():
                 graph_only_radio = gr.Radio(choices=[True, False], value=True, label="Graph-only Answer")
                 graph_vector_radio = gr.Radio(choices=[True, False], value=False, label="Graph-Vector Answer")
+            with gr.Column():
+                with gr.Row():
+                    extraction_method_input = gr.Radio(choices=["LLM", "Hybrid", "TextRank"],
+                                                       value="Hybrid", label="Keywords Extraction Method")
+                    max_keyword_num = gr.Number(value=5, label="Max Keywords Num", precision=5)
 
             def toggle_slider(enable):
                 return gr.update(interactive=enable)
@@ -322,6 +357,9 @@ def toggle_slider(enable):
             custom_related_information,
             answer_prompt_input,
             keywords_extract_prompt_input,
+            extraction_method_input,
+            mask_words_input,
+            max_keyword_num,
             example_num,
         ],
         outputs=[raw_out, vector_only_out, graph_only_out, graph_vector_out],
@@ -387,6 +425,9 @@ def several_rag_answer(
         custom_related_information_ui: str,
         answer_prompt: str,
         keywords_extract_prompt: str,
+        keywords_extraction_method: str,
+        mask_words: str,
+        keyword_num: int,
         answer_max_line_count_ui: int = 1,
         progress=gr.Progress(track_tqdm=True),
     ):
@@ -406,6 +447,9 @@ def several_rag_answer(
                 custom_related_information_ui,
                 answer_prompt,
                 keywords_extract_prompt,
+                keywords_extraction_method,
+                mask_words,
+                keyword_num,
             )
             df.at[index, "Basic LLM Answer"] = basic_llm_answer
             df.at[index, "Vector-only Answer"] = vector_only_answer
@@ -439,10 +483,13 @@ def several_rag_answer(
             custom_related_information,
             answer_prompt_input,
             keywords_extract_prompt_input,
+            extraction_method_input,
+            mask_words_input,
+            max_keyword_num,
             answer_max_line_count,
         ],
         outputs=[qa_dataframe, gr.File(label="Download Answered File", min_width=40)],
     )
     questions_file.change(read_file_to_excel, questions_file, [qa_dataframe, answer_max_line_count])
     answer_max_line_count.change(change_showing_excel, answer_max_line_count, qa_dataframe)
-    return inp, answer_prompt_input, keywords_extract_prompt_input, custom_related_information
+    return inp, answer_prompt_input, keywords_extract_prompt_input, custom_related_information, mask_words_input

hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py

@@ -52,6 +52,23 @@ def stopwords(self, lang: str = "chinese") -> List[str]:
 
         return self._stopwords[lang]
 
+    def check_nltk_data(self):
+        required_packages = ['punkt', 'punkt_tab', 'averaged_perceptron_tagger', "averaged_perceptron_tagger_eng"]
+        cache_dir = self.get_cache_dir()
+        nltk_data_dir = os.environ.get("NLTK_DATA", cache_dir)
+
+        if nltk_data_dir not in nltk.data.path:
+            nltk.data.path.append(nltk_data_dir)
+
+        for package in required_packages:
+            try:
+                if package in ['punkt', 'punkt_tab']:
+                    nltk.data.find(f'tokenizers/{package}')
+                else:
+                    nltk.data.find(f'taggers/{package}')
+            except LookupError:
+                nltk.download(package, download_dir=nltk_data_dir)
+
     @staticmethod
     def get_cache_dir() -> str:
         """Locate a platform-appropriate cache directory for hugegraph-llm,