# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
import os
from pathlib import Path

import torch
import torch.distributed as dist
from optimum.onnxruntime import ORTModelForCausalLM
from transformers import AutoConfig, AutoTokenizer, GenerationConfig

device_id = 0
device = torch.device(f"cuda:{device_id}")  # Change to torch.device("cpu") if running on CPU

ep = "CUDAExecutionProvider"  # change to CPUExecutionProvider if running on CPU
ep_options = {"device_id": device_id}

model_id = "mistralai/Mistral-7B-Instruct-v0.2"
model_path = "./Olive/examples/llama2/models/qlora/qlora-conversion-transformers_optimization-bnb_quantization/gpu-cuda_model"

model_path = Path(model_path)

if not (model_path / "config.json").exists():
    config = AutoConfig.from_pretrained(model_id)
    config.save_pretrained(model_path)
else:
    config = AutoConfig.from_pretrained(model_path)

if not (model_path / "generation_config.json").exists():
    gen_config = GenerationConfig.from_pretrained(model_id)
    gen_config.save_pretrained(model_path)
else:
    gen_config = GenerationConfig.from_pretrained(model_path)

tokenizer = AutoTokenizer.from_pretrained(model_id)

model = ORTModelForCausalLM.from_pretrained(
    model_path,
    config=config,
    generation_config=gen_config,
    use_io_binding=True,
    # provider="CUDAExecutionProvider",
    provider=ep,
    provider_options={"device_id": device_id}
    # provider_options={"device_id": str(rank)},
)