Buckets:
| """Load the merged text decision model; no PEFT adapters or base download needed.""" | |
| import importlib | |
| import json | |
| from pathlib import Path | |
| import torch | |
| from transformers import AutoTokenizer | |
| class Head256(torch.nn.Module): | |
| def __init__(self, hidden): | |
| super().__init__() | |
| self.register_buffer('mu', torch.zeros(1, hidden)) | |
| self.register_buffer('sd', torch.ones(1, hidden)) | |
| self.linear = torch.nn.Linear(hidden, 256, dtype=torch.float32) | |
| def forward(self, features, counts): | |
| z = self.linear((features.float() - self.mu) / self.sd) | |
| return z.masked_fill(torch.arange(256, device=z.device)[None] >= counts[:, None], -1e30) | |
| def load_model(path, device='cuda'): | |
| torch.set_float32_matmul_precision('highest') | |
| torch.backends.cuda.matmul.allow_tf32 = False | |
| torch.backends.cudnn.allow_tf32 = False | |
| path = Path(path) | |
| c = json.loads((path / 'decision_config.json').read_text()) | |
| assert c['backbone_module'].startswith('transformers.models.') | |
| cls = getattr(importlib.import_module(c['backbone_module']), c['backbone_class']) | |
| model = cls.from_pretrained(path / 'backbone', dtype=torch.float32, | |
| attn_implementation='sdpa').to(device).eval() | |
| buffers = torch.load(path / 'runtime_buffers.pt', map_location='cpu', weights_only=True) | |
| for name, value in buffers.items(): | |
| parent, _, attr = name.rpartition('.') | |
| module = model.get_submodule(parent) if parent else model | |
| setattr(module, attr, value.to(device)) | |
| head = Head256(c['hidden_size']).to(device).eval() | |
| head.load_state_dict(torch.load(path / 'head.pt', map_location=device, weights_only=True)) | |
| tok = AutoTokenizer.from_pretrained(path) | |
| assert tok.chat_template | |
| return model, head, tok | |
| def predict(model, head, tokenizer, state, question, options): | |
| assert 1 <= len(options) <= 256 | |
| choices = '\n'.join(f'{i+1}. {v}' for i,v in enumerate(options)) | |
| prompt = (f'{state}\n\n---\n\nQUESTION: {question}\n\nOPTIONS:\n{choices}\n\n' | |
| f'Reply with only the number of the correct option (1-{len(options)}).\n' | |
| 'Output a single number and nothing else.') | |
| msgs = [{'role':'user', 'content':prompt}] | |
| ids = None | |
| for kw in ({'enable_thinking':False}, {}): | |
| try: | |
| ids = tokenizer.apply_chat_template(msgs, add_generation_prompt=True, tokenize=True, **kw) | |
| break | |
| except (TypeError, ValueError): | |
| msgs = [{'role':'user', 'content':[{'type':'text', 'text':prompt}]}] | |
| try: | |
| ids = tokenizer.apply_chat_template(msgs, add_generation_prompt=True, tokenize=True, **kw) | |
| break | |
| except (TypeError, ValueError): | |
| continue | |
| if ids is None: raise ValueError('Could not format prompt') | |
| if hasattr(ids, 'keys'): ids=ids['input_ids'] | |
| if ids and isinstance(ids[0], list): ids=ids[0] | |
| device = next(model.parameters()).device | |
| ids = torch.tensor([ids], device=device) | |
| with torch.autocast('cuda', dtype=torch.bfloat16): | |
| out = model(input_ids=ids, attention_mask=torch.ones_like(ids), use_cache=False) | |
| hidden = out.last_hidden_state if hasattr(out, 'last_hidden_state') else out[0] | |
| z = head(hidden[:, -1].float(), torch.tensor([len(options)], device=device)) | |
| return z.softmax(-1)[0, :len(options)].cpu().tolist() | |
Xet Storage Details
- Size:
- 3.42 kB
- Xet hash:
- 7d5a1ee86d94ff7b333a6c9f119caafe1307d87d3943aa064fe25f25c6129d0d
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.