import torchimport torch.nn as nnimport torch.nn.functional as Ffrom transformers import AutoTokenizer, AutoModelForCausalLM DEVICE = "cuda" if torch.cuda.is_available() else "cpu" class QwenTokenizerAdapter: def __init__(self, model_id): self.hf = AutoTokenizer.from_pretrained(model_id) self.eos_id = self.hf.eos_token_id self.bos_id = self.hf.bos_token_id if self.hf.bos_token_id is not None else self.eos_id self.pad_id = self.hf.pad_token_id if self.hf.pad_token_id is not None else self.eos_id self.vocab_size = len(self.hf) self.tokenizer = self def encode(self, text, add_special_tokens=False): return self.hf.encode(text, add_special_tokens=add_special_tokens) def decode(self, ids, skip_special_tokens=False): return self.hf.decode(ids, skip_special_tokens=skip_special_tokens) def id_to_token(self, idx): return self.hf.convert_ids_to_tokens([int(idx)])[0] def tokens(self, text): return self.hf.tokenize(text) class HFModelWrapper(nn.Module): """Matches TinyGPT's call convention -- model(xb, yb) -> (logits, loss), model.generate(idx, max_new_tokens=, eos_id=, top_k=) -> full sequence.""" def __init__(self, hf_model): super().__init__() self.hf_model = hf_model def forward(self, xb, yb=None): logits = self.hf_model(input_ids=xb).logits loss = None if yb is not None: loss = F.cross_entropy( logits.view(-1, logits.size(-1)), yb.view(-1), ignore_index=-100, ) return logits, loss def generate(self, idx, max_new_tokens, eos_id=None, top_k=None, temperature=None, do_sample=None): return self.hf_model.generate( input_ids=idx, max_new_tokens=max_new_tokens, eos_token_id=eos_id, pad_token_id=eos_id, top_k=top_k, do_sample=do_sample if do_sample is not None else (top_k is not None or temperature is not None), temperature=temperature if temperature is not None else 1.0, ) @property def vocab_size(self): return self.hf_model.config.vocab_size def load_hf_checkpoint(path, device): hf_model = AutoModelForCausalLM.from_pretrained( path, torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, attn_implementation="sdpa", ).to(device) return HFModelWrapper(hf_model).to(device) tokenizer = QwenTokenizerAdapter(REPO)model = load_hf_checkpoint(REPO, DEVICE).eval()