from __future__ import annotations
import json
import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoModelForCausalLM, AutoTokenizer
from tinycenn_lm.smollm2_memory_fusion import SmolMemoryFusionConfig, replace_attention_layers

def load_model(repo_id: str, token=None, device=None):
    device=torch.device(device or ('cuda' if torch.cuda.is_available() else 'cpu'))
    dtype=torch.bfloat16 if device.type=='cuda' and torch.cuda.is_bf16_supported() else (torch.float16 if device.type=='cuda' else torch.float32)
    meta_path=hf_hub_download(repo_id,'tinycenn_model.json',token=token)
    with open(meta_path,encoding='utf-8') as f:
        meta=json.load(f)
    tokenizer=AutoTokenizer.from_pretrained(repo_id,token=token,use_fast=True)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token=tokenizer.eos_token
    model=AutoModelForCausalLM.from_pretrained(meta['base_model'],dtype=dtype)
    cfg=SmolMemoryFusionConfig.from_dict(meta['memory_fusion'])
    replace_attention_layers(model,cfg,meta['accepted_layers'])
    state_path=hf_hub_download(repo_id,'model.safetensors',token=token)
    state=load_file(state_path,device='cpu')
    model.load_state_dict(state,strict=True)
    model.to(device).eval(); model.requires_grad_(False); model.config.use_cache=False
    if hasattr(model,'generation_config'):
        model.generation_config.use_cache=False
    return model,tokenizer,meta