Export to HuggingFace

import asyncio
import torch
from pathlib import Path
from huggingface_hub import HfApi, HfFolder
import mint
from mint import types
from transformers import AutoModelForCausalLM, AutoTokenizer

async def train_and_publish_to_hub():
    service_client = mint.ServiceClient()
    
    # Step 1: Train 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)
    
    # Train on instruction tuning examples
    training_examples = [
        "Write a haiku about spring:\nGreen leaves emerge soft\nWarmth returns to sleeping earth\nLife renews again",
        "Explain quantum computing:\nQuantum computers process information using qubits in superposition.",
    ]
    
    for i, example in enumerate(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()
        print(f"  Example {i+1}: trained")
    
    # Step 2: Save checkpoint
    checkpoint = await training_client.save_weights_for_sampler_async(
        name="poetry-assistant-v1"
    )
    checkpoint = await checkpoint.result_async()
    print("Checkpoint saved")
    
    # Step 3: Download and merge weights
    print("\n=== Downloading and Merging Weights ===")
    
    base_model_id = "Qwen/Qwen3-0.6B"
    base_model = AutoModelForCausalLM.from_pretrained(
        base_model_id,
        torch_dtype=torch.bfloat16,
        device_map="auto",
    )
    base_tokenizer = AutoTokenizer.from_pretrained(base_model_id)
    
    # Download LoRA weights (use MinT API)
    # lora_weights = checkpoint.download_weights(...)
    # merged_model = merge_lora_weights(base_model, lora_weights)
    
    # For demo, use base model
    merged_model = base_model
    
    # Step 4: Create model directory with metadata
    print("\n=== Preparing for Hub Upload ===")
    
    model_dir = Path("./poetry-assistant")
    model_dir.mkdir(exist_ok=True)
    
    # Save model and tokenizer
    merged_model.save_pretrained(model_dir)
    base_tokenizer.save_pretrained(model_dir)
    
    # Create a README
    readme_content = """
    # Poetry Assistant
    
    Fine-tuned Qwen3-0.6B for poetry and creative writing.
    
    ## Training
    
    - Base model: `Qwen/Qwen3-0.6B`
    - Training method: LoRA fine-tuning on MinT
    - Rank: 16
    - Learning rate: 5e-5
    - Training examples: 2
    
    ## Usage
    
    ```python
    from transformers import AutoModelForCausalLM, AutoTokenizer
    
    model = AutoModelForCausalLM.from_pretrained("your-username/poetry-assistant")
    tokenizer = AutoTokenizer.from_pretrained("your-username/poetry-assistant")
    
    prompt = "Write a haiku about spring:"
    inputs = tokenizer.encode(prompt, return_tensors="pt")
    outputs = model.generate(inputs, max_length=50)
    print(tokenizer.decode(outputs[0]))
    ```
    
    ## License
    
    Same as base model.
    """
    
    with open(model_dir / "README.md", "w") as f:
        f.write(readme_content)
    
    # Step 5: Upload to HuggingFace Hub
    print("\n=== Uploading to HuggingFace Hub ===")
    
    # Authenticate (assumes `huggingface-cli login` was run)
    # Or set HUGGING_FACE_HUB_TOKEN environment variable
    hf_api = HfApi()
    
    hub_repo_id = "your-username/poetry-assistant"
    
    # Create repo if it doesn't exist
    try:
        hf_api.create_repo(repo_id=hub_repo_id, private=False)
        print(f"Created repo: {hub_repo_id}")
    except Exception as e:
        print(f"Repo already exists or error: {e}")
    
    # Upload files
    hf_api.upload_folder(
        folder_path=str(model_dir),
        repo_id=hub_repo_id,
        commit_message="Initial release: LoRA fine-tuned on MinT",
    )
    
    print(f"Model uploaded to: https://huggingface.co/{hub_repo_id}")
    print("\nYour model is now public and can be loaded with:")
    print(f'  model = AutoModelForCausalLM.from_pretrained("{hub_repo_id}")')

asyncio.run(train_and_publish_to_hub())