refactor

Author: JiahangXu

examples/search_r1/search_r1_agent.py

@@ -4,10 +4,9 @@
 
 import os
 import re
-import shutil
-import tempfile
 import time
-from typing import Any, Dict, List, Literal, Optional, Tuple, TypedDict, cast
+import pandas as pd
+from typing import Any, Dict, List, Optional, Tuple, TypedDict, cast
 
 
 import requests
@@ -98,14 +97,10 @@ def passages2string(retrieval_result: List[RetrievalItem]) -> str:
 def call_llm(
     llm_client: OpenAI,
     model_name: str,
-    content: str = "",
     messages: List[dict] = [],
     temperature: float = 1.0,
     max_tokens: int = 500,
 ) -> str:
-    if not len(messages):
-        messages=[{"role": "user", "content": content}]
-    print(messages)
     response = llm_client.chat.completions.create(
         model=model_name,
         messages=messages,

examples/search_r1/train_search_r1_agent.py

@@ -76,7 +76,7 @@
         "experiment_name": "searchr1_minibatch256_runner32",
         "nnodes": 1,
         "test_freq": 10,
-        "save_freq":10,
+        "save_freq": 10,
         "total_epochs": 15,
         "total_training_steps": 300,
         "default_local_dir": "/mnt/teamdrive/search_r1/searchr1_checkpoints/searchr1_minibatch256_runner32/"