Publish to MinT Hub

import asyncio
import mint
from mint import types

async def train_and_publish_to_hub():
    service_client = mint.ServiceClient()
    
    # Step 1：在 MinT 上训练 model
    print("=== Training on MinT ===")
    
    training_client = await service_client.create_lora_training_client_async(
        base_model="Qwen/Qwen3-0.6B",
        rank=16,
    )
    tokenizer = training_client.get_tokenizer()
    adam_params = types.AdamParams(learning_rate=5e-5)
    
    # 简单的训练循环
    training_examples = [
        "Customer service training example 1",
        "Customer service training example 2",
    ]
    
    for example in training_examples:
        tokens = tokenizer.encode(example)
        model_input = types.ModelInput.from_ints(tokens[:-1])
        target_tokens = tokens[1:]
        weights = [1.0] * len(target_tokens)
        
        datum = types.Datum(
            model_input=model_input,
            loss_fn_inputs={"target_tokens": target_tokens, "weights": weights},
        )
        
        result = await training_client.forward_backward_async([datum], loss_fn="cross_entropy")
        await result.result_async()
        
        optim_future = training_client.optim_step_async(adam_params)
        await optim_future.result_async()
    
    # Step 2：把 checkpoint 发布到 MinT Hub
    print("\n=== Publishing to MinT Hub ===")
    
    checkpoint_name = "customer-service-v1"
    checkpoint = await training_client.save_weights_for_sampler_async(
        name=checkpoint_name,
    )
    checkpoint = await checkpoint.result_async()
    
    # 带元数据发布到 Hub
    hub_metadata = {
        "description": "Fine-tuned for customer service conversations",
        "tags": ["customer-service", "qwen3", "lora"],
        "base_model": "Qwen/Qwen3-0.6B",
        "rank": 16,
        "training_examples": len(training_examples),
    }
    
    # 发布（MinT API 调用）
    try:
        rest_client = service_client.get_rest_client()
        rest_client.publish_checkpoint(
            checkpoint_name=checkpoint_name,
            hub_id="customer-service/qwen-v1",
            metadata=hub_metadata,
        )
        print(f"Published to MinT Hub: customer-service/qwen-v1")
    except Exception as e:
        print(f"Note: MinT Hub publishing requires server support: {e}")
    
    # Step 3：从 Hub 加载用于推理
    print("\n=== Loading from MinT Hub ===")
    
    # 其它用户可以这样加载：
    try:
        sampling_client = service_client.create_sampling_client_from_hub(
            hub_id="customer-service/qwen-v1",
        ).result()
        print("Loaded model from MinT Hub")
        
        # 生成
        prompt_ids = tokenizer.encode("Customer: I have a problem")
        prompt = types.ModelInput.from_ints(prompt_ids)
        
        output = sampling_client.sample(
            prompt,
            sampling_params=types.SamplingParams(max_tokens=64, temperature=0.7),
        ).result()
        
        response = tokenizer.decode(output.sequences[0].tokens)
        print(f"Response: {response}")
    except Exception as e:
        print(f"Note: Load from Hub requires server support: {e}")
    
    # Step 4：列出已发布的 model
    print("\n=== Hub Model Management ===")
    
    try:
        models = rest_client.list_hub_models(org="customer-service")
        for model in models:
            print(f"  - {model['name']} (v{model['version']}, {model['downloads']} downloads)")
    except Exception as e:
        print(f"Note: Hub listing requires server support: {e}")
    
    # Step 5：设置 model 元数据（便于检索）
    try:
        rest_client.update_hub_model_metadata(
            hub_id="customer-service/qwen-v1",
            metadata={
                "readme": "# Customer Service Assistant\n\nFine-tuned for handling customer inquiries...",
                "license": "CC-BY-4.0",
                "authors": ["your-team"],
            }
        )
        print("Updated model metadata")
    except Exception as e:
        print(f"Note: Metadata updates require server support: {e}")

asyncio.run(train_and_publish_to_hub())