jasinliu commented on code in PR #73:
URL:
https://github.com/apache/incubator-hugegraph-ai/pull/73#discussion_r1732255982
##########
hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py:
##########
@@ -16,59 +16,119 @@
# under the License.
-from typing import Dict, Any, List, Literal
+from typing import Literal, Dict, Any, List, Optional, Tuple
import jieba
from hugegraph_llm.models.embeddings.base import BaseEmbedding
+from hugegraph_llm.models.rerankers.init_reranker import Rerankers
from nltk.translate.bleu_score import sentence_bleu
-def get_score(query: str, content: str) -> float:
+def get_bleu_score(query: str, content: str) -> float:
query_tokens = jieba.lcut(query)
content_tokens = jieba.lcut(content)
return sentence_bleu([query_tokens], content_tokens)
class MergeDedupRerank:
def __init__(
- self,
- embedding: BaseEmbedding,
- topk: int = 10,
- strategy: Literal["bleu", "priority"] = "bleu"
+ self,
+ embedding: BaseEmbedding,
+ topk: int = 20,
+ graph_ratio: float = 0.5,
+ method: Literal["bleu", "reranker"] = "bleu",
+ near_neighbor_first: bool = False,
+ custom_related_information: Optional[str] = None,
):
+ assert method in [
+ "bleu",
+ "reranker",
+ ], f"Unimplemented rerank method '{method}'."
self.embedding = embedding
+ self.graph_ratio = graph_ratio
self.topk = topk
- if strategy == "bleu":
- self.rerank_func = self._bleu_rerank
- elif strategy == "priority":
Review Comment:
> Do you mean that this priority parameter applies to all reranking methods?
If so, it makes sense.
Is this ok?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
