import torch
import os

def pack_for_prod(input_path, output_path):
    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 any(x in name for x in ["q_proj", "k_proj", "v_proj", "out_proj", "ffn_w"]):
            print(f"  -> Пакуем слой: {name}")
            # Твоя формула квантования
            gamma = weight.abs().mean().clamp(min=1e-9)
            w_quant = torch.round(weight / gamma).clamp(-1, 1).to(torch.int8)
            
            new_sd[name] = w_quant
            # СОХРАНЯЕМ ГАММУ С ТОЧНЫМ ИМЕНЕМ
            gamma_key = name.replace(".weight", "_gamma")
            new_sd[gamma_key] = gamma.half()
        else:
            new_sd[name] = weight.half()
            
    torch.save(new_sd, output_path)
    print(f"✅ Готово! Файл сохранен в {output_path}")

pack_for_prod("jiarck_pro_1b_model.pt", "jirack_pro_1b_prod.pt")