| import torch |
|
|
| from lm_eval.api.model import LM |
| from lm_eval.api.registry import register_model |
|
|
| import sys |
| sys.path.insert(0, "runtime") |
|
|
| from litgpt import Config |
| from litgpt.model import GPT |
|
|
|
|
| @register_model("obsidian_multiscreen") |
| class ObsidianMultiscreenLM(LM): |
|
|
| def __init__( |
| self, |
| checkpoint, |
| device="cuda", |
| dtype="bfloat16", |
| backend="triton", |
| **kwargs |
| ): |
|
|
| self.device = device |
|
|
| if backend: |
| import os |
| os.environ["MULTISCREEN_BACKEND"] = backend |
|
|
| config = Config.from_file( |
| f"{checkpoint}/model_config.yaml" |
| ) |
|
|
| self.model = GPT(config) |
|
|
| state = torch.load( |
| f"{checkpoint}/lit_model.pth", |
| map_location="cpu" |
| ) |
|
|
| self.model.load_state_dict( |
| state["model"] |
| ) |
|
|
| self.model.to(device) |
|
|
| if dtype == "bfloat16": |
| self.model.to(torch.bfloat16) |
|
|
| self.model.eval() |
|
|
| self.vocab_size = config.padded_vocab_size |
| self.max_length = config.block_size |
|
|
|
|
| @property |
| def eot_token_id(self): |
| return 0 |
|
|
|
|
| @property |
| def max_length(self): |
| return self._max_length |
|
|
|
|
| @max_length.setter |
| def max_length(self, x): |
| self._max_length=x |
|
|
|
|
| def tok_encode(self, string): |
| return self.tokenizer.encode(string) |
|
|
|
|
| def loglikelihood(self, requests): |
|
|
| results=[] |
|
|
| for request in requests: |
|
|
| context, continuation = request.args |
|
|
| text=context+continuation |
|
|
| ids=torch.tensor( |
| [self.tokenizer.encode(text)], |
| device=self.device |
| ) |
|
|
| with torch.no_grad(): |
| logits=self.model(ids) |
|
|
| log_probs=torch.log_softmax( |
| logits, |
| dim=-1 |
| ) |
|
|
| cont_ids=self.tokenizer.encode( |
| continuation |
| ) |
|
|
| score=0 |
|
|
| for i,tok in enumerate(cont_ids): |
| score += log_probs[ |
| 0, |
| -(len(cont_ids)-i+1), |
| tok |
| ] |
|
|
| results.append( |
| (float(score), True) |
| ) |
|
|
| return results |
