f"""你是一个逻辑推理专家，擅长解决逻辑推理问题。以下是一个逻辑推理的题目，形式为单项选择题。所有的问题都是（close-world assumption）闭世界假设，即未观测事实都为假。请逐步分析问题并在最后一行输出答案，最后一行的格式为"答案是：A"。

### 题目:
{problem}

### 问题:
{question}
{options}
"""

f"""你是一个逻辑推理专家，擅长解决逻辑推理问题。以下是一个逻辑推理的题目，形式为单项选择题。所有的问题都是（close-world assumption）闭世界假设，即未观测事实都为假。请认真思考，逐步分析问题并在最后一行输出答案，最后一行输出的格式为"答案是：<>"。
### 题目:
{problem}

### 问题:
{question}
{options}
***Instructions***
请根据上述题目和问题，逐步思考推理过程，并最终输出答案。

1. 问题形式为单项选择题，选项中存在且仅存在一个正确答案。
2. 所有的问题都是（close-world assumption）闭世界假设，即未观测事实都为假。
3. 逐步输出思考过程和解答推理过程，过程尽可能简要精炼。
4. 输出为：答案是：<>。
"""

f"""你是一个逻辑推理专家，擅长解决逻辑推理问题。以下是一个逻辑推理的题目：
### 题目
{problem}

### 问题
{question}
{options}

***Instructions***
请根据上述题目，以同样的格式生成类似的题目和相对应的选项，注意：

1. 基于同样的题目,生成不同的问题和选项,选项中必须有且只有一个正确答案。
2. 每次只生成一个问题,形式为单项选择题，确保选项中存在且仅存在一个正确答案。
3. 所有的问题都是（close-world assumption）闭世界假设，即未观测事实都为假。
4. 请认真思考，逐步分析问题并在最后一行输出答案。
5. 输出格式如下：
### 问题
生成的问题
生成的选项
答案

***Example***
### 问题
{question}
{options}
{answer}
"""

swift sft \
    --model_type qwen2_5-32b-instruct \
    --model_id_or_path /root/autodl-tmp/Qwen2.5-32B-Instruct \
    --model_revision master \
    --sft_type lora \
    --template_type AUTO \
    --dtype bf16 \
    --output_dir /root/autodl-tmp/qwen \
    --dataset $TRAIN_DATA \
    --num_train_epochs 3 \
    --max_length 4096 \
    --check_dataset_strategy warning \
    --lora_target_modules ALL \
    --gradient_checkpointing true \
    --lora_rank 256 \
    --lora_alpha 512 \
    --batch_size 1 \
    --learning_rate 1e-4 \
    --gradient_accumulation_steps 8 \
    --max_grad_norm 1.0 \
    --warmup_ratio 0.03 \
    --save_steps 2000 \
    --save_total_limit 3 \
    --use_flash_attn true \
    --save_only_model true \
    --logging_steps 100

swift export \
    --model_type qwen2_5-32b-instruct \
    --model_id_or_path /root/autodl-tmp/Qwen2.5-32B-Instruct \
    --ckpt_dir $OUTPUT_DIR  \
    --dtype fp16 \
    --merge_device_map auto \
    --merge_lora true

swift export \
    --model_type qwen2_5-32b-instruct \
    --quant_bits 4 \
    --quant_method gptq \
    --model_id_or_path $OUTPUT_DIR-merged \
    --dataset $TRAIN_DATA alpaca-zh alpaca-en sharegpt-gpt4:default

from swift.llm import (
    ModelType, get_vllm_engine, get_default_template_type,
    get_template, inference_vllm, inference_stream_vllm
)
def chat_swift(texts, model_path,
               model_type='qwen2_5_32b_instruct',
               max_tokens=2048,
               max_model_len=2048,
               temperature=0.1,
               top_p=0.7,
               top_k=20, 
               use_tqdm=True,
               **kwargs):
    model_type = getattr(ModelType, model_type)
    llm_engine = get_vllm_engine(model_type, model_id_or_path=model_path, **kwargs)
    template_type = get_default_template_type(model_type)
    template = get_template(template_type, llm_engine.hf_tokenizer)
    llm_engine.generation_config.max_new_tokens = max_tokens
    llm_engine.generation_config.temperature = temperature
    llm_engine.generation_config.top_p = top_p
    llm_engine.generation_config.top_k = top_k
    generation_info = {}

    request_list = [{'query': x} for x in texts]
    outputs = inference_vllm(llm_engine, template, request_list, generation_info=generation_info, use_tqdm=use_tqdm)
    print(generation_info)

    outputs = [out['response'] for out in outputs]
    return outputs

# vllm调用
# prompts为构造好的测试集问题prompt
# model_path 为量化后的模型路径
outputs = chat_swift(prompts, model_path, quantization='gptq')

def ensemble(answers):
    # answers 为四个模型的预测答案，第0位为最优模型的答案
    a = answers[0]
    b = np.unique(answers[1:])
    if len(b) == 1:
        return b[0]
    return a