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: Train a model on MinT
    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)
    
    # Quick training loop
    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: Save checkpoint to 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()
    
    # Publish to Hub with metadata
    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),
    }
    
    # Publish (MinT API call)
    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: Load from Hub for inference
    print("\n=== Loading from MinT Hub ===")
    
    # Users can now load with:
    try:
        sampling_client = service_client.create_sampling_client_from_hub(
            hub_id="customer-service/qwen-v1",
        ).result()
        print("Loaded model from MinT Hub")
        
        # Generate
        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: List published models
    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: Set model metadata (for discovery)
    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())