import torch
import os

def pack_for_prod(input_path, output_path):
    if not os.path.exists(input_path):
        print(f"❌ Исходный файл {input_path} не найден!")
        return

    print(f"📦 Загружаем исходник: {input_path}")
    sd = torch.load(input_path, map_location="cpu")
    if "model_state_dict" in sd: 
        sd = sd["model_state_dict"]

    new_sd = {}
    for name, weight in sd.items():
        if not isinstance(weight, torch.Tensor):
            continue

        # Убираем возможные префиксы параллелизма (на всякий случай)
        clean_name = name.replace("_orig_mod.", "").replace("module.", "")

        # Если это веса слоев, которые должны быть тернарными/INT8
        if any(x in clean_name for x in ["q_proj", "k_proj", "v_proj", "out_proj", "ffn_w"]):
            print(f"  -> Пакуем слой: {clean_name}")
            
            # ВНИМАНИЕ: Если модель УЖЕ обучалась как тернарная, 
            # НЕ применяй формулу квантования заново! Просто сделай .to(torch.int8).
            # Оставляю формулу ТОЛЬКО если ты уверен, что делаешь PTW-квантование:
            if weight.dtype != torch.int8:
                gamma = weight.abs().mean().clamp(min=1e-9)
                w_quant = torch.round(weight / gamma).clamp(-1, 1).to(torch.int8)
                
                new_sd[clean_name] = w_quant.contiguous()
                
                # Имя гаммы должно СТРОГО соответствовать твоей архитектуре из fine_tune скрипта!
                # Если в архитектуре скейл называется по-другому, замени "_gamma" ниже:
                gamma_key = clean_name.replace(".weight", "_gamma")
                new_sd[gamma_key] = gamma.half().contiguous()
            else:
                # Если веса уже были int8, просто сохраняем их как есть
                new_sd[clean_name] = weight.contiguous()
        else:
            # Эмбеддинги, нормы и lm_head переводим во float16 для прод-инференса
            new_sd[clean_name] = weight.half().contiguous()

    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    torch.save(new_sd, output_path)
    print(f"✅ Готово! Файл сохранен в {output_path}")

if __name__ == "__main__":
    # Исправил опечатку в путях на правильные (из твоих прошлых запросов)
    INPUT = "/mnt/nfs_share/JiRackTernaryPro_1b/jiarck_pro_1b_model.pt"
    OUTPUT = "/mnt/nfs_share/JiRackTernaryPro_1b/jirack_pro_1b_prod.pt"
    
    pack_for_prod(INPUT, OUTPUT)