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：训练 LoRA model（或加载已有 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()
    
    # 简单训练（demo 只跑 1 步）
    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：保存 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：下载 LoRA weights
    print("\n=== Downloading LoRA Weights ===")
    
    # 给下载下来的 weights 准备目录
    weights_dir = Path("./downloaded_lora")
    weights_dir.mkdir(exist_ok=True)
    
    # 实操中通过 MinT API 或 REST endpoint 下载 weights
    # 这里是占位
    # lora_weights = checkpoint.download_weights(save_path=weights_dir)
    # print(f"Weights downloaded to {weights_dir}")
    
    # Step 4：和 base model 合并并保存
    print("\n=== Merging LoRA with Base Model ===")
    
    # 从 HuggingFace 加载 base model 和 tokenizer
    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)
    
    # 合并 LoRA weights（实操中用 peft 库）
    # 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)
    
    # demo 里直接用 base model
    merged_model = base_model
    
    # Step 5：把合并后的 model 保存到本地
    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：本地推理
    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)
    
    # 生成
    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())