1from pathlib import Path
2import torch
3from huggingface_hub import snapshot_download
4from texformer.models.checkpoint_loader import load_texformer_model
5from texformer.tokenization.tokenizer import TeXFormerTokenizer
6
7repo_id = "aamingem/texformer-576m-bf16"
8local_dir = Path(snapshot_download(repo_id=repo_id))
9tokenizer_dir = local_dir / "tokenizer"
10
11tokenizer = TeXFormerTokenizer(tokenizer_dir)
12if torch.cuda.is_available():
13 device = torch.device("cuda")
14elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
15 device = torch.device("mps")
16else:
17 device = torch.device("cpu")
18model, global_step, epoch, missing, unexpected = load_texformer_model(
19 checkpoint_path=local_dir / "model.pt",
20 tokenizer=tokenizer,
21 device=device,
22)
23print("Loaded model:", model.__class__.__name__)
24print("Missing keys:", len(missing), "Unexpected keys:", len(unexpected))
25