Upload RAIF/COMPLEX_INSTRUCTIONS with huggingface_hub

RAIF/COMPLEX_INSTRUCTIONS

@@ -0,0 +1,462 @@
+import os, json, requests, re, traceback, sys
+from queue import Queue, Empty
+import argparse
+import threading
+import asyncio
+import aiohttp
+import time
+import pandas as pd
+import xlsxwriter
+from tqdm import tqdm
+from datetime import datetime
+from importlib import reload
+import logging
+import concurrent.futures
+from transformers import AutoTokenizer
+
+sys.path.insert(0, "/mnt")
+from disk2.api.deepseek_v3_2 import deepseek_v3_2
+from disk2.api.deepseek_v3_2_thinking import deepseek_v3_2_thinking
+from disk2.api.deepseek_v4_pro import deepseek_v4_pro
+from disk2.api.deepseek_v4_flash import deepseek_v4_flash
+from disk2.api.qwen3_235b_a22b import qwen3_235b_a22b
+from disk2.api.gpt51 import gpt51
+from disk2.api.vllm_server import vllm_server
+from disk2.api.openai_server import openai_server
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--input_path", type=str, default="")
+parser.add_argument("--input_file_type", type=str, default="jsonl")
+parser.add_argument("--save_path", type=str, default="")
+parser.add_argument("--model_id", type=str, default="")
+parser.add_argument("--model_url", type=str, default="")
+
+parser.add_argument("--question_type", type=str, default="conversations") #! TODO
+parser.add_argument("--save_freq", type=int, default=10)
+parser.add_argument("--qkey", type=str, default="q") # 输入文件种query对应的key
+parser.add_argument("--akey", type=str, default="default") # 输出文件中answer对应的key
+parser.add_argument("--system_prompt", type=str, default="You are a helpful assistant.")
+parser.add_argument("--max_tokens", type=int, default=2048)
+parser.add_argument("--resume", action="store_true", default=False)
+parser.add_argument("--api_key", type=str, default="xxx")
+parser.add_argument("--response_format", type=str, default="")
+parser.add_argument("--verbose", type=bool, default=True)
+parser.add_argument("--maxtry", type=int, default=3)
+parser.add_argument("--tensor_parallel_size", type=int, default=4)
+
+
+
+args = parser.parse_args()
+global client
+
+
+def process_streaming(response):
+    start_time = time.time()
+    collected_chunks = []
+    collected_messages = []
+    for chunk in response:
+        chunk_time = time.time() - start_time  # calculate the time delay of the chunk
+        collected_chunks.append(chunk)  # save the event response
+        chunk_message = chunk.choices[0].delta.content  # extract the message
+        collected_messages.append(chunk_message)  # save the message
+    print(f"Full response received {chunk_time:.2f} seconds after request")
+    collected_messages = [m for m in collected_messages if m is not None]
+    full_reply_content = "".join(collected_messages)
+    return full_reply_content
+
+
+
+def build_model():
+    model_name_or_path = "/mnt/disk2/models/Qwen2.5-7B-Instruct"
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)
+    print("loaded tokenizer qwen")
+    return tokenizer
+
+
+
+def build_model_mistral():
+    model_name_or_path = "{YOUR_PATH_TO_PRETRAINED_MODELS}/pretrained_models/Mistral-7B-Instruct-v0.3"
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)
+    print("loaded tokenizer mistral")
+    return tokenizer
+
+
+
+def build_model_llama():
+    model_name_or_path = "{YOUR_PATH_TO_PRETRAINED_MODELS}/pretrained_models/Meta-Llama-3.1-8B-Instruct_meta-llama"
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)
+    print("loaded tokenizer llama")
+    return tokenizer
+
+
+
+def count_token(response, tokenizer):
+    inputs = tokenizer.encode(response)
+    return len(inputs)
+
+
+
+def count_token_max(response, tokenizer_list):
+    response_token_len_list = [count_token(response, tokenizer_item) for tokenizer_item in tokenizer_list]
+    return max(response_token_len_list)
+
+
+
+# global tokenizer
+# if "mistral" in (args.model_id).lower() or "ministral" in (args.model_id).lower():
+#     tokenizer = build_model_mistral()
+# elif "llama" in (args.model_id).lower():
+#     tokenizer = build_model_llama()
+# else:
+#     tokenizer = build_model()
+
+
+if args.question_type == "conversations":
+    args.qkey = "conversations"
+
+
+def prepare_batch_item(index, data):
+    if "tools" in data:
+        tools = json.loads(data["tools"])
+    else:
+        tools = None
+
+    if args.question_type == "conversations":
+        messages = data[args.qkey]
+    else:
+        messages = [
+            {
+                "role":"system",
+                "content":args.system_prompt,
+            },
+            {
+                "role":"user",
+                "content":data[args.qkey],
+            }
+        ]
+        
+
+    messages_new = []
+
+    for message in messages:
+        role = message["role"]
+        content = message["content"]
+
+        messages_new.append(
+            {
+                "role":role,
+                "content":content,
+            }
+        )
+    while messages_new[-1]["role"] == "assistant":
+        messages_new.pop(-1)
+
+    messages = messages_new
+    messages_str = ""
+    for message in messages:
+        if message["content"] is not None:
+            messages_str += message["content"]+"\n"
+
+        if "tool_calls" in message and message["tool_calls"] is not None:
+            messages_str += str(message["tools"])+"\n"
+
+    return {
+        "index": index,
+        "data": data,
+        "messages": messages,
+        "messages_str": messages_str,
+        "tools": tools,
+    }
+
+
+
+def finalize_batch_item(item):
+    data = item["data"]
+    response = {
+        "choices": [
+            {
+                "message": {
+                    "content": item["response"].get("content") if 'response' in item else None,
+                    "reasoning_content": item["response"].get("reasoning_content") if 'response' in item else None,
+                    "tool_calls": item["response"].get("tool_calls") if 'response' in item else None,
+                }
+            }
+        ]
+    }
+
+    result = response["choices"][0]["message"]["content"]
+    # assert (result is not None) or ("tool_calls" in response["choices"][0]["message"] and response["choices"][0]["message"]["tool_calls"] is not None)
+
+    if result and result.startswith("<answer>"):
+        result = (result[len("<answer>"):]).lstrip()
+    if "reasoning_content" in response["choices"][0]["message"]:
+        reasoning = response["choices"][0]["message"]["reasoning_content"]
+    else:
+        reasoning = None
+
+    if "tool_calls" in response["choices"][0]["message"]:
+        tool_calls = response["choices"][0]["message"]["tool_calls"]
+    else:
+        tool_calls = None
+
+    if args.akey == "default":
+        data[args.model_id] = result
+        if reasoning:
+            data[args.model_id + "_reasoning"] = reasoning
+        if tool_calls:
+            data[args.model_id + "_tool_calls"] = tool_calls
+    else:
+        data[args.akey] = result
+        if reasoning:
+            data[args.akey + "_reasoning"] = reasoning
+        if tool_calls:
+            data[args.model_id + "_tool_calls"] = tool_calls
+
+    return item["index"], data
+
+
+
+def save_results(results, save_path):
+    """
+    分批次保存输出结果, a+ mode
+    results: [index, result]
+    """
+    with open(save_path, "a+") as f:
+        for results_i in results:
+            f.writelines(json.dumps(results_i[1], ensure_ascii=False)+"\n")
+
+
+def load_input_file(input_path, file_type="jsonl"):
+    if file_type == "jsonl":
+        with open(args.input_path) as f:
+            Data = f.readlines()
+        Data = [json.loads(i) for i in Data]
+    elif file_type == "json":
+        with open(args.input_path) as f:
+            Data = json.load(f)
+    else:
+        raise ValueError("file_type must be jsonl or json")
+    if args.question_type == "conversations":
+        for i in range(len(Data)):
+            #! 删掉最后的assistant
+            if Data[i][args.qkey][-1]["role"] == "assistant":
+                Data[i][args.qkey].pop(-1)
+
+    return Data
+
+
+def get_resume_state(save_path, file_type="jsonl"):
+    count = 0
+    # 如果文件不存在
+    if not os.path.exists(save_path):
+        return count
+    # 如果文件存在，加载处理进度
+    if file_type == "jsonl":
+        with open(save_path, "r", encoding="utf-8") as f:
+            for line in f:
+                if line:
+                    count += 1
+    elif file_type == "json":
+        with open(save_path, "r", encoding="utf-8") as f:
+            count = len(json.load(f))
+    else:
+        raise ValueError("file_type must be jsonl or json")
+    return count
+    
+
+def main():
+    global client
+
+    print('================================================')
+    print(f"[INFO] Starting inference with model_id: {args.model_id}, model_url: {args.model_url}")
+    print(f"[INFO] Input path: {args.input_path}, Input file type: {args.input_file_type}")
+    print(f"[INFO] Save path: {args.save_path}")
+    print(f"[INFO] Question type: {args.question_type}, Qkey: {args.qkey}, Answer key: {args.akey}")
+    print('================================================')
+
+
+    os.makedirs(os.path.dirname(args.save_path), exist_ok=True)
+    if args.resume:
+        resume_state = get_resume_state(args.save_path, args.input_file_type)
+    else:
+        resume_state = 0
+
+    # 加载输入文件
+    Data = load_input_file(args.input_path, args.input_file_type)
+    print(f"Data loaded. Total {len(Data)} records. {len(Data)-resume_state} records to run.")
+    # cut Data
+    Data = Data[resume_state:]
+
+    # 加载模型服务地址
+    model_url = args.model_url
+    model_id = args.model_id
+    print("加载模型服务地址：", model_url)
+    print("加载模型名称：", args.model_id)
+    max_threads = 200
+    llm = None
+    sampling_params = None
+    if model_url not in [None, "", "None"]:
+        from vllm import LLM, SamplingParams
+
+        print("加载本地离线 vLLM 模型：", model_url)
+        llm = LLM(
+            model=model_url,
+            gpu_memory_utilization=0.9,
+            tensor_parallel_size=args.tensor_parallel_size,
+        )
+        sampling_params = SamplingParams(max_tokens=args.max_tokens)
+        # llm = LLM(
+        #     model=model_url,
+        #     gpu_memory_utilization=0.9,
+        #     tensor_parallel_size=1,
+        #     pipeline_parallel_size=8,
+        #     max_num_seqs=1000
+        # )
+        print(
+            f"Initialized offline vLLM with model={model_url}, "
+            f"max_tokens={args.max_tokens}"
+        )
+        max_threads = 1024
+    elif model_id == "deepseek_v3_2":
+        response_format = args.response_format if args.response_format != "" else "text"
+        client = deepseek_v3_2(response_format=response_format, cache_dir="./cache/deepseek_v3_2_cache", think_enabled='disabled')
+    elif model_id == "deepseek_v3_2_thinking":
+        response_format = args.response_format if args.response_format != "" else "text"
+        client = deepseek_v3_2_thinking(response_format=response_format, cache_dir="./cache/deepseek_v3_2_thinking_1_cache", think_enabled='enabled')
+    elif model_id == "deepseek_v4_pro":
+        response_format = args.response_format if args.response_format != "" else "text"
+        client = deepseek_v4_pro(response_format=response_format, cache_dir="./cache/deepseek_v4_pro_cache", think_enabled='enabled', reasoning_effort='high')
+    elif model_id == "deepseek_v4_flash":
+        response_format = args.response_format if args.response_format != "" else "text"
+        client = deepseek_v4_flash(response_format=response_format, cache_dir="./cache/deepseek_v4_flash_cache", think_enabled='enabled', reasoning_effort='high')
+    elif model_id == "qwen3_235b_a22b":
+        response_format = args.response_format if args.response_format != "" else "text"
+        client = qwen3_235b_a22b(response_format=response_format, cache_dir="./cache/qwen3_235b_a22b_cache")
+    elif model_id == "gpt51":
+        response_format = args.response_format if args.response_format != "" else "text"
+        client = gpt51(response_format=response_format, cache_dir="./cache/gpt51_cache")
+    elif model_id == "openai_server_crab":
+        client = openai_server(cache_dir="./models/cache/openai_server_cache")        
+    elif model_id == "openai_server_conifer":
+        client = openai_server(cache_dir="./models/cache/openai_server_conifer_2_cache") 
+    elif model_id == "openai_server_ultraif":
+        client = openai_server(cache_dir="./models/cache/openai_server_ultraif_cache") 
+    elif model_id == "openai_server_llama3_70b":
+        client = openai_server(cache_dir="./models/cache/openai_server_llama3_70b_cache") 
+    elif model_id == "openai_server_llama3_crab":
+        client = openai_server(cache_dir="./models/cache/openai_server_llama3_crab_cache") 
+    else:
+        assert False, "Unsupported model_id, please provide a valid model_id or model_url."
+
+    
+    
+    continuous_results = [] # 待保存的连续输出结果序列
+    tbar = tqdm(total=len(Data)+resume_state, initial=resume_state, desc="Processing")
+    batch_items = [prepare_batch_item(i, Data[i]) for i in range(len(Data))]
+
+    if args.verbose:
+        for item in batch_items[:5]:
+            print('[INFO] 传递给模型的消息内容如下：')
+            print(json.dumps(item["messages"], ensure_ascii=False, indent=4))
+            print('================================================')
+
+
+    failed_items = list(batch_items)
+    num_retry = 0
+    while failed_items and num_retry < 1:
+
+
+        # for item in failed_items:
+        #     print(item['messages'])
+        #     input()
+
+
+
+        if llm is not None:
+            message_batches = [item["messages"] for item in failed_items]
+            outputs = []
+            offline_batch_size = 1000
+            for start_idx in range(0, len(message_batches), offline_batch_size):
+                message_batch = message_batches[start_idx:start_idx + offline_batch_size]
+                vllm_outputs = llm.chat(
+                    message_batch,
+                    sampling_params,
+                    use_tqdm=True,
+                    chat_template_kwargs={"enable_thinking": True},
+                )
+                for output in vllm_outputs:
+                    # print(output)
+                    raw_generated_text = None
+                    finish_reason = None
+                    generated_text = None
+                    reasoning_content = None
+                    if output.outputs:
+                        raw_generated_text = output.outputs[0].text
+                        finish_reason = getattr(output.outputs[0], "finish_reason", None)
+                        if '</think>' in raw_generated_text:
+                            generated_text = raw_generated_text.split('</think>')[-1].strip()
+                            reasoning_content = raw_generated_text.split('</think>')[0].replace('<think>', '').strip()
+                            # print(generated_text)
+                            # print(reasoning_content)
+                            # assert False
+                        else:
+                            generated_text = raw_generated_text.strip()
+                    outputs.append(
+                        {
+                            "content": generated_text,
+                            "reasoning_content": reasoning_content,
+                            "raw_response": {
+                                "prompt": output.prompt,
+                                "finish_reason": finish_reason,
+                            },
+                        }
+                    )
+        else:
+            outputs = client(
+                [item["messages"] for item in failed_items],
+                use_cache=(num_retry == 0),
+                max_threads=max_threads,
+                maxtry=args.maxtry,
+            )
+        
+        # if args.verbose:
+        #     for item, output in zip(failed_items[:3], outputs[:3]):
+        #         print('[INFO] 模型的输出内容如下：')
+        #         print(json.dumps(output, ensure_ascii=False, indent=4))
+        #         print('================================================')
+        
+
+        next_failed_items = []
+        for item, output in zip(failed_items, outputs):
+            if output is None:
+                next_failed_items.append(item)
+                continue
+            item["response"] = output
+
+        if next_failed_items:
+            print("!"*50)
+            num_retry += 1
+            time.sleep(5)
+            failed_items = next_failed_items
+            continue
+
+        failed_items = []
+
+    for item in batch_items:
+        index, result = finalize_batch_item(item)
+        continuous_results.append([index, result])
+        tbar.update(1)
+
+        if len(continuous_results) >= args.save_freq:
+            save_results(continuous_results, args.save_path)
+            continuous_results = []
+
+    # 保存剩余的结果
+    if continuous_results:
+        save_results(continuous_results, args.save_path)
+    return
+
+
+if __name__ == "__main__":
+    main()