from trio import ServiceClient, SampleInput, AdamParams, SamplingParams
import time
import os
from datasets import load_dataset

model_id = "Qwen/Qwen2.5-3B"
batch_size = 2

domain = os.getenv("SERVICE_DOMAIN", "117.50.44.26:3000")

client = ServiceClient(domain)
print("支持的模型列表：", client.get_supported_models())

train_client = client.create_lora_training_client("Qwen/Qwen2.5-3B", train_unembed=False)
tokenizer = train_client.get_tokenizer()

# ----------------------- 数据集处理 -----------------------

dataset = load_dataset("HuggingFaceTB/smoltalk", "everyday-conversations", split="train[:10]")

all_samples = []

for example in dataset:
    # 将对话消息应用模板 (Chat Template)
    # 注意：这里假设你想对整个对话进行训练
    text = tokenizer.apply_chat_template(example["messages"], tokenize=False)
    tokens = tokenizer.encode(text, add_special_tokens=False)

    # 构造输入
    model_input = tokens[:-1]
    target_tokens = tokens[1:]
    weights = [1] * len(target_tokens)

    # 构造 SampleInput
    sample = SampleInput(
        model_input=list(model_input),
        loss_fn_inputs={
            "weights": weights,
            "target_tokens": list(target_tokens)
        }
    )
    all_samples.append(sample)

# 将列表按 batch_size 切分
batches = [all_samples[i: i + batch_size] for i in range(0, len(all_samples), batch_size)]

# ----------------------- 训练循环 -----------------------

for i, batch in enumerate(batches):
    # 前向 + 反向
    start_fb = time.perf_counter()
    res = train_client.forward_backward(batch)
    print('forward_backward: ', res)

    end_fb = time.perf_counter()
    print(f'------- {i} 前向+反向执行时间：{end_fb - start_fb} s')

    # 梯度优化
    start_opt = time.perf_counter()

    adam_params = AdamParams(learning_rate=0.001, beta1=0.9, beta2=0.999, eps=1e-8, weight_decay=0)
    res = train_client.optim_step(adam_params)
    print('optimize_step: ', res)

    end_opt = time.perf_counter()
    print(f'------- {i} 梯度优化执行时间：{end_opt - start_opt} s')

# ----------------------- 保存模型权重 -----------------------
start_w = time.perf_counter()
save_res = train_client.save_weights_for_sampler("train001")
print(save_res)
end_w = time.perf_counter()

print(f'------- 保存模型权重执行时间：{end_w - start_w} s')

# ----------------------- Lora 模型推理 -----------------------

sample_client = train_client.create_sampling_client(model_path=save_res['path'])
input_ids = tokenizer.encode("how are you？")

# sample 根据模型生成文本补全
params = SamplingParams(max_tokens=50, temperature=0.7)
res = sample_client.sample(prompt=[input_ids], num_samples=1, sampling_params=params)
print(res)
