up

README.md

@@ -9,8 +9,150 @@ base_model:
 
 # DeepRefine-v1-8B
 
-A general LLM-based reasoning model for agent-compiled knowledge refinement that improves the quality of any pre-constructed knowledge bases with user queries to make it more9suitable for the downstream tasks.
+A general LLM-based reasoning model for agent-compiled knowledge refinement that improves the quality of any pre-constructed knowledge bases with user queries to make it more suitable for the downstream tasks.
 
-## Quick Start
+## Simple Demo
+Below is a single refinement process with respect to a specfic query and pre-defined results from its interactions with any databases.
 
-Todo.
+```python
+DEEPREFINE_JUDGEMENT_SYSTEM_PROMPT = """
+As an advanced judgement assistant, your task is to judge whether the given question is answerable based on the provided KG context.
+
+Evaluate whether the given question is answerable based on the provided KG context. Output your judgment in the following format:
+<judge>Yes</judge> or <judge>No</judge>
+
+**Important:** You must think carefully about the question and the KG context before making your judgment. And output your judgment result directly in the specified format.
+"""
+
+DEEPREFINE_JUDGEMENT_USER_PROMPT = """
+Question: {question}
+Knowledge Graph (KG) context: {triples_string}
+"""
+
+DEEPREFINE_ERROR_ABDUCTION_SYSTEM_PROMPT = """
+As an advanced error abduction assistant, your task is to analyze the error reasons based on the given interaction history.
+
+Analyze the reasons of the unanswerable questions based on the given interaction history from the incompleteness, incorrectness, and redundancy perspectives. Output your analysis in the following format:
+<abduction>...</abduction>
+
+**Important:** You must think carefully about the interaction history before making your analysis. And output your analysis result directly in the specified format.
+"""
+
+DEEPREFINE_ERROR_ABDUCTION_USER_PROMPT = """
+Interaction history: {interaction_history}
+"""
+
+DEEPREFINE_ACTION_SYSTEM_PROMPT = """
+As an advanced knowledge graph refinement assistant, your task is to generate a series of actions (**within 10 actions**) to refine the given KG to make it more suitable for answering the given question.
+
+Based on the given KG and the analysed error reasons, refine the given KG to make it more easily for retrieval and answering the given question. You have the following three types of actions to conduct:
+
+- insert_edge(subject, relation, object): Insert a new edge into the KG to complete the missing information.
+- delete_edge(subject, relation, object): Delete an edge from the KG to remove the redundant information or conflicting information.
+- replace_node(old_entity, new_entity): Replace an entity in the KG to correct the errors or deal with disambiguation.
+
+Output a series of actions (**within 10 actions**) in the following format:
+<refinement>insert_edge("...", "...", "...")|delete_edge("...", "...", "...")|replace_node("...", "...")|...</refinement>
+
+**Important:** You must think carefully about the given KG and the analysed error reasons before making your refinement. DO NOT DELETE ANY IRRELEVANT TRIPLES FROM THE ORIGINAL KG. TRY TO KEEP THE ORIGINAL KG AS MUCH AS POSSIBLE. DO NOT GENERATE TOO MANY ACTIONS. And output your refinement result directly in the specified format.
+"""
+
+DEEPREFINE_ACTION_USER_PROMPT = """
+Original Text: {original_text}
+KG: {triples_string}
+Question: {question}
+Error reasons: {error_reasons}
+"""
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+model_name = "HaoyuHuang/DeepRefine-v1-8B"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    device_map="auto" if device == "cuda" else None,
+    trust_remote_code=True,
+)
+if device != "cuda":
+    model = model.to(device)
+
+def call_model(system_prompt: str, user_prompt: str, max_new_tokens: int = 8192) -> str:
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_prompt},
+    ]
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(text, return_tensors="pt").to(model.device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=False,
+            temperature=0.0,
+        )
+
+    generated_ids = outputs[0, inputs["input_ids"].shape[-1]:]
+    return tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
+
+@dataclass
+class RetrievalStepResult:
+    """
+    For single step inference result, for debugging / analysis.
+    """
+    num_hops: int
+    base_top_k: int
+    query: str
+    retrieved_subgraph: List[Dict[str, str]]
+    raw_response: str
+    answerable: bool
+    answer: Optional[str] = None
+
+question = "your_question"
+triples_string = "retrieved_triples:
+interaction_history = "list_of_RetrievalStepResult"
+original_text = "triple_related_original_text"
+
+## Answerability Judgement Call
+judgement_user_prompt = DEEPREFINE_JUDGEMENT_USER_PROMPT.format(
+    question=question,
+    triples_string=triples_string,
+)
+judgement_result = call_model(
+    DEEPREFINE_JUDGEMENT_SYSTEM_PROMPT,
+    judgement_user_prompt
+)
+print("Answerability Judgement:", judgement_result)
+
+
+## Error Abduction Call
+abduction_user_prompt = DEEPREFINE_ERROR_ABDUCTION_USER_PROMPT.format(
+    interaction_history=interaction_history,
+)
+abduction_result = call_model(
+    DEEPREFINE_ERROR_ABDUCTION_SYSTEM_PROMPT,
+    abduction_user_prompt,
+)
+print("Error Abduction:", abduction_result)
+
+
+## Refinement Actions Generation Call
+actions_user_prompt = DEEPREFINE_ACTION_USER_PROMPT.format(
+    original_text=original_text,
+    triples_string=triples_string,
+    question=question,
+    error_reasons=abduction_result,
+)
+actions_result = call_model(
+    DEEPREFINE_ACTION_SYSTEM_PROMPT,
+    actions_user_prompt,
+)
+print("Refinement Actions:", actions_result)
+
+## parse the refinement actions and make an interface with the operators of your own databases
+# ...
+```