From e888d9f2ea6d11816952f271145be1ff670fa867 Mon Sep 17 00:00:00 2001 From: chenzihong-gavin Date: Tue, 2 Dec 2025 19:03:33 +0800 Subject: [PATCH 1/2] fix: fix bds baseline --- baselines/BDS/bds.py | 25 +++++++++---------------- 1 file changed, 9 insertions(+), 16 deletions(-) diff --git a/baselines/BDS/bds.py b/baselines/BDS/bds.py index 33c36718..2f64e00f 100644 --- a/baselines/BDS/bds.py +++ b/baselines/BDS/bds.py @@ -1,15 +1,15 @@ import argparse import asyncio import json -import os -from dataclasses import dataclass from typing import List import networkx as nx from dotenv import load_dotenv from tqdm.asyncio import tqdm as tqdm_async -from graphgen.models import NetworkXStorage, OpenAIClient, Tokenizer +from graphgen.bases import BaseLLMWrapper +from graphgen.models import NetworkXStorage +from graphgen.operators import init_llm from graphgen.utils import create_event_loop QA_GENERATION_PROMPT = """ @@ -52,10 +52,12 @@ def _post_process(text: str) -> dict: return {} -@dataclass class BDS: - llm_client: OpenAIClient = None - max_concurrent: int = 1000 + def __init__(self, synthesizer_llm_client: BaseLLMWrapper = None): + self.llm_client: BaseLLMWrapper = synthesizer_llm_client or init_llm( + "synthesizer" + ) + self.max_concurrent: int = 1000 def generate(self, tasks: List[dict]) -> List[dict]: loop = create_event_loop() @@ -102,16 +104,7 @@ async def job(item): load_dotenv() - tokenizer_instance: Tokenizer = Tokenizer( - model_name=os.getenv("TOKENIZER_MODEL", "cl100k_base") - ) - llm_client = OpenAIClient( - model_name=os.getenv("SYNTHESIZER_MODEL"), - api_key=os.getenv("SYNTHESIZER_API_KEY"), - base_url=os.getenv("SYNTHESIZER_BASE_URL"), - tokenizer_instance=tokenizer_instance, - ) - bds = BDS(llm_client=llm_client) + bds = BDS() graph = NetworkXStorage.load_nx_graph(args.input_file) From a5944e940c9c6f9260c86e0cd6ba6567875f9646 Mon Sep 17 00:00:00 2001 From: chenzihong <58508660+ChenZiHong-Gavin@users.noreply.github.com> Date: Tue, 2 Dec 2025 19:13:57 +0800 Subject: [PATCH 2/2] Update baselines/BDS/bds.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- baselines/BDS/bds.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/baselines/BDS/bds.py b/baselines/BDS/bds.py index 2f64e00f..419fdcc3 100644 --- a/baselines/BDS/bds.py +++ b/baselines/BDS/bds.py @@ -53,11 +53,11 @@ def _post_process(text: str) -> dict: class BDS: - def __init__(self, synthesizer_llm_client: BaseLLMWrapper = None): - self.llm_client: BaseLLMWrapper = synthesizer_llm_client or init_llm( + def __init__(self, llm_client: BaseLLMWrapper = None, max_concurrent: int = 1000): + self.llm_client: BaseLLMWrapper = llm_client or init_llm( "synthesizer" ) - self.max_concurrent: int = 1000 + self.max_concurrent: int = max_concurrent def generate(self, tasks: List[dict]) -> List[dict]: loop = create_event_loop()