From b2aecf833bb203efde350957ddceb90747c37577 Mon Sep 17 00:00:00 2001 From: Boris Arzentar Date: Sat, 20 Apr 2024 20:16:33 +0200 Subject: [PATCH] fix: add max_tokens to all ExtractKnowledgeGraph calls --- cognee/modules/cognify/evaluate.py | 3 ++- .../knowledge_graph/extract_knowledge_graph_module.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/cognee/modules/cognify/evaluate.py b/cognee/modules/cognify/evaluate.py index adc7ae634..fe4588352 100644 --- a/cognee/modules/cognify/evaluate.py +++ b/cognee/modules/cognify/evaluate.py @@ -36,7 +36,8 @@ def evaluate(): evaluate_on_hotpotqa = Evaluate(devset = devset, num_threads = 1, display_progress = True, display_table = 5, max_tokens = 4096) - compiled_extract_knowledge_graph = ExtractKnowledgeGraph() + gpt4 = dspy.OpenAI(model = config.openai_model, api_key = config.openai_key, model_type = "chat", max_tokens = 4096) + compiled_extract_knowledge_graph = ExtractKnowledgeGraph(lm = gpt4) compiled_extract_knowledge_graph.load(get_absolute_path("./programs/extract_knowledge_graph/extract_knowledge_graph.json")) def evaluate_answer(example, graph_prediction, trace = None): diff --git a/cognee/modules/data/extraction/knowledge_graph/extract_knowledge_graph_module.py b/cognee/modules/data/extraction/knowledge_graph/extract_knowledge_graph_module.py index 85c95efab..04ff42c8a 100644 --- a/cognee/modules/data/extraction/knowledge_graph/extract_knowledge_graph_module.py +++ b/cognee/modules/data/extraction/knowledge_graph/extract_knowledge_graph_module.py @@ -41,7 +41,7 @@ def are_all_nodes_connected(graph: KnowledgeGraph) -> bool: class ExtractKnowledgeGraph(dspy.Module): - def __init__(self, lm = dspy.OpenAI(model = config.openai_model, api_key = config.openai_key, model_type = "chat")): + def __init__(self, lm = dspy.OpenAI(model = config.openai_model, api_key = config.openai_key, model_type = "chat", max_tokens = 4096)): super().__init__() self.lm = lm dspy.settings.configure(lm=self.lm)