LoRA Adapter

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

async def lora_merge_and_deploy():
    service_client = mint.ServiceClient()
    
    # Step 1: Train a LoRA model (or load an existing checkpoint)
    print("=== Training LoRA Model ===")
    
    training_client = await service_client.create_lora_training_client_async(
        base_model="Qwen/Qwen3-0.6B",
        rank=16,
    )
    tokenizer = training_client.get_tokenizer()
    
    # Quick training (1 step for demo)
    text = "Hello, world! This is a test."
    tokens = tokenizer.encode(text)
    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()
    
    adam_params = types.AdamParams(learning_rate=5e-5)
    await training_client.optim_step_async(adam_params).result_async()
    
    # Step 2: Save the checkpoint
    checkpoint = await training_client.save_weights_for_sampler_async(
        name="lora-checkpoint-v1"
    )
    checkpoint = await checkpoint.result_async()
    print(f"Checkpoint saved: lora-checkpoint-v1")
    
    # Step 3: Download LoRA weights
    print("\n=== Downloading LoRA Weights ===")
    
    # Create a directory for downloaded weights
    weights_dir = Path("./downloaded_lora")
    weights_dir.mkdir(exist_ok=True)
    
    # In practice, use the MinT API or REST endpoint to download weights
    # For now, we simulate this step
    # lora_weights = checkpoint.download_weights(save_path=weights_dir)
    # print(f"Weights downloaded to {weights_dir}")
    
    # Step 4: Merge with base model and save
    print("\n=== Merging LoRA with Base Model ===")
    
    # Load base model and tokenizer from HuggingFace
    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)
    
    # Merge LoRA weights (in practice, use peft library)
    # from peft import get_peft_model_state_dict
    # lora_state = torch.load(weights_dir / "adapter_model.bin")
    # merged_model = merge_lora_weights(base_model, lora_state)
    
    # For demo, we just use the base model
    merged_model = base_model
    
    # Step 5: Save merged model locally
    output_dir = Path("./merged_model")
    output_dir.mkdir(exist_ok=True)
    
    merged_model.save_pretrained(output_dir)
    base_tokenizer.save_pretrained(output_dir)
    print(f"Merged model saved to {output_dir}")
    
    # Step 6: Local inference
    print("\n=== Local Inference ===")
    
    local_model = AutoModelForCausalLM.from_pretrained(
        output_dir,
        torch_dtype=torch.bfloat16,
        device_map="auto",
    )
    local_tokenizer = AutoTokenizer.from_pretrained(output_dir)
    
    # Generate
    prompt = "The capital of France is"
    inputs = local_tokenizer.encode(prompt, return_tensors="pt").to(local_model.device)
    
    with torch.no_grad():
        outputs = local_model.generate(
            inputs,
            max_length=50,
            temperature=0.7,
            top_p=0.9,
        )
    
    generated_text = local_tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(f"Generated: {generated_text}")
    
    print("\n=== Deployment Complete ===")
    print(f"Merged model is ready for deployment at {output_dir}")
    print("You can now:")
    print("  - Serve with vLLM: vllm serve ./merged_model")
    print("  - Deploy to mobile with ONNX conversion")
    print("  - Upload to HuggingFace Hub")

asyncio.run(lora_merge_and_deploy())