Merge pull request #95 from CHERRY-ui8/update-baselines-and-clients

fix: OpenAIClient parameter from model_name to model

Author: ChenZiHong-Gavin

baselines/Genie/genie.py

@@ -122,7 +122,7 @@ async def process_chunk(content: str):
     load_dotenv()
 
     llm_client = OpenAIClient(
-        model_name=os.getenv("SYNTHESIZER_MODEL"),
+        model=os.getenv("SYNTHESIZER_MODEL"),
         api_key=os.getenv("SYNTHESIZER_API_KEY"),
         base_url=os.getenv("SYNTHESIZER_BASE_URL"),
     )

baselines/LongForm/longform.py

@@ -89,7 +89,7 @@ async def process_chunk(content: str):
     load_dotenv()
 
     llm_client = OpenAIClient(
-        model_name=os.getenv("SYNTHESIZER_MODEL"),
+        model=os.getenv("SYNTHESIZER_MODEL"),
         api_key=os.getenv("SYNTHESIZER_API_KEY"),
         base_url=os.getenv("SYNTHESIZER_BASE_URL"),
     )

baselines/SELF-QA/self-qa.py

@@ -156,7 +156,7 @@ async def process_chunk(content: str):
     load_dotenv()
 
     llm_client = OpenAIClient(
-        model_name=os.getenv("SYNTHESIZER_MODEL"),
+        model=os.getenv("SYNTHESIZER_MODEL"),
         api_key=os.getenv("SYNTHESIZER_API_KEY"),
         base_url=os.getenv("SYNTHESIZER_BASE_URL"),
     )

baselines/Wrap/wrap.py

@@ -109,7 +109,7 @@ async def process_chunk(content: str):
     load_dotenv()
 
     llm_client = OpenAIClient(
-        model_name=os.getenv("SYNTHESIZER_MODEL"),
+        model=os.getenv("SYNTHESIZER_MODEL"),
         api_key=os.getenv("SYNTHESIZER_API_KEY"),
         base_url=os.getenv("SYNTHESIZER_BASE_URL"),
     )

graphgen/models/llm/api/openai_client.py

@@ -32,7 +32,7 @@ class OpenAIClient(BaseLLMWrapper):
     def __init__(
         self,
         *,
-        model_name: str = "gpt-4o-mini",
+        model: str = "gpt-4o-mini",
         api_key: Optional[str] = None,
         base_url: Optional[str] = None,
         json_mode: bool = False,
@@ -44,7 +44,7 @@ def __init__(
         **kwargs: Any,
     ):
         super().__init__(**kwargs)
-        self.model_name = model_name
+        self.model = model
         self.api_key = api_key
         self.base_url = base_url
         self.json_mode = json_mode
@@ -109,7 +109,7 @@ async def generate_topk_per_token(
         kwargs["max_tokens"] = 1
 
         completion = await self.client.chat.completions.create(  # pylint: disable=E1125
-            model=self.model_name, **kwargs
+            model=self.model, **kwargs
         )
 
         tokens = get_top_response_tokens(completion)
@@ -141,7 +141,7 @@ async def generate_answer(
             await self.tpm.wait(estimated_tokens, silent=True)
 
         completion = await self.client.chat.completions.create(  # pylint: disable=E1125
-            model=self.model_name, **kwargs
+            model=self.model, **kwargs
         )
         if hasattr(completion, "usage"):
             self.token_usage.append(

webui/app.py

@@ -42,7 +42,7 @@ def init_graph_gen(config: dict, env: dict) -> GraphGen:
 
     tokenizer_instance = Tokenizer(config.get("tokenizer", "cl100k_base"))
     synthesizer_llm_client = OpenAIClient(
-        model_name=env.get("SYNTHESIZER_MODEL", ""),
+        model=env.get("SYNTHESIZER_MODEL", ""),
         base_url=env.get("SYNTHESIZER_BASE_URL", ""),
         api_key=env.get("SYNTHESIZER_API_KEY", ""),
         request_limit=True,
@@ -51,7 +51,7 @@ def init_graph_gen(config: dict, env: dict) -> GraphGen:
         tokenizer=tokenizer_instance,
     )
     trainee_llm_client = OpenAIClient(
-        model_name=env.get("TRAINEE_MODEL", ""),
+        model=env.get("TRAINEE_MODEL", ""),
         base_url=env.get("TRAINEE_BASE_URL", ""),
         api_key=env.get("TRAINEE_API_KEY", ""),
         request_limit=True,