1from pathlib import Path
2import torch
3from huggingface_hub import snapshot_download
4from texformer.models.model import TeXFormer, TeXFormerConfig
5from texformer.models.quantize import dequantize_weight_only_state_dict
6from texformer.tokenization.tokenizer import TeXFormerTokenizer
7
8repo_id = "aamingem/texformer-270m-int8_weight_only"
9local_dir = Path(snapshot_download(repo_id=repo_id))
10tokenizer_dir = local_dir / "tokenizer"
11checkpoint = torch.load(local_dir / "model.pt", map_location="cpu", weights_only=False)
12
13if torch.cuda.is_available():
14 device = torch.device("cuda")
15elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
16 device = torch.device("mps")
17else:
18 device = torch.device("cpu")
19
20config = TeXFormerConfig(**checkpoint["config"])
21state_dict = dequantize_weight_only_state_dict(checkpoint["quantized_state_dict"])
22model = TeXFormer(config)
23model.load_state_dict(state_dict, strict=False)
24model = model.to(device)
25
26tokenizer = TeXFormerTokenizer(tokenizer_dir)
27print("Loaded dequantized model with tokenizer:", tokenizer.pdf_vocab_size, tokenizer.latex_vocab_size)
28