import os
import torch
from transformers import AutoTokenizer
from sllm_store.transformers import load_model, save_model

# model_name = "facebook/opt-1.3b"
model_name = "Qwen/Qwen2.5-1.5b"
model_folder = os.getenv("MODEL_FOLDER") # change this to your own path
model_path = os.path.join(model_folder, model_name)
# =======================================================================================================================
torch.cuda.empty_cache()
# warm up the GPU
num_gpus = torch.cuda.device_count()
for i in range(num_gpus):
    torch.ones(1).to(f"cuda:{i}")
    torch.cuda.synchronize()

model = load_model(
    model_name,
    device_map="auto",
    storage_path=model_folder, 
    fully_parallel=True,
)
# =======================================================================================================================
tokenizer = AutoTokenizer.from_pretrained(model_name)
inputs = tokenizer("Hello, my dog is cute", return_tensors="pt").to("cuda")
outputs = model.generate(**inputs)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

for name, param in model.named_parameters():
    print(f"{name}: shape={param.shape}, dtype={param.dtype}")

import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# model_name = "facebook/opt-1.3b"
model_name = "Qwen/Qwen2.5-1.5b"
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
model = model.to('cuda')
# =======================================================================================================================
num_gpus = torch.cuda.device_count()
for i in range(num_gpus):
    torch.ones(1).to(f"cuda:{i}")
    torch.cuda.synchronize()

tokenizer = AutoTokenizer.from_pretrained(model_name)
inputs = tokenizer("Hello, my dog is cute", return_tensors="pt").to("cuda")
outputs = model.generate(**inputs)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

for name, param in model.named_parameters():
    print(f"{name}: shape={param.shape}, dtype={param.dtype}")