import json
import os

from trio import Datum, ServiceClient

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

servicer = ServiceClient(domain)

trainer = servicer.create_lora_training_client(
    base_model="Qwen/Qwen2.5-3B",
)

sample = Datum(
    model_input=[1, 2, 3, 4],
    loss_fn_inputs={
        "target_tokens": [2, 3, 4, 5],
        "logprobs": [-0.2, -0.3, -0.4, -0.5],
        "advantages": [0.4, 0.3, 0.2, 0.1],
    },
)

for epoch in range(1):
    result = trainer.forward_backward(
        data=[sample, sample],
        loss_fn="importance_sampling",
    )
    assert isinstance(result, dict)

    print(json.dumps(result, indent=2, ensure_ascii=False))
