A class used to perform inference with the BitNetTransformer model. ... Attributes ---------- model : torch.nn.Module an instance of the BitNetTransformer model device : str the device to run the model on ('cpu' or 'cuda') Methods ------- load_
| 6 | |
| 7 | |
| 8 | class BitNetInference: |
| 9 | """ |
| 10 | A class used to perform inference with the BitNetTransformer model. |
| 11 | |
| 12 | ... |
| 13 | |
| 14 | Attributes |
| 15 | ---------- |
| 16 | model : torch.nn.Module |
| 17 | an instance of the BitNetTransformer model |
| 18 | device : str |
| 19 | the device to run the model on ('cpu' or 'cuda') |
| 20 | |
| 21 | Methods |
| 22 | ------- |
| 23 | load_model(model_path) |
| 24 | Loads a trained model from a .pth file. |
| 25 | generate(input_str, length) |
| 26 | Generates a sequence of tokens based on the input string. |
| 27 | """ |
| 28 | |
| 29 | def __init__(self, device="cuda"): |
| 30 | """ |
| 31 | Parameters |
| 32 | ---------- |
| 33 | device : str, optional |
| 34 | The device to run the model on ('cpu' or 'cuda'). By default, 'cuda' is used. |
| 35 | """ |
| 36 | self.device = device |
| 37 | self.model = BitNetTransformer(num_tokens=256, dim=512, depth=8) |
| 38 | self.model = AutoregressiveWrapper(self.model, max_seq_len=1024) |
| 39 | self.model.to(self.device) |
| 40 | |
| 41 | def load_model(self, model_path): |
| 42 | """Loads a trained model from a .pth file.""" |
| 43 | self.model.load_state_dict(torch.load(model_path, weights_only=True)) |
| 44 | self.model.eval() |
| 45 | |
| 46 | @staticmethod |
| 47 | def decode_token(token): |
| 48 | """Decodes a token into a character.""" |
| 49 | return str(chr(max(32, token))) |
| 50 | |
| 51 | @staticmethod |
| 52 | def decode_tokens(tokens): |
| 53 | """Decodes a sequence of tokens into a string.""" |
| 54 | return "".join(list(map(BitNetInference.decode_token, tokens))) |
| 55 | |
| 56 | def generate(self, input_str, length): |
| 57 | """Generates a sequence of tokens based on the input string.""" |
| 58 | inp = ( |
| 59 | torch.from_numpy(np.fromstring(input_str, dtype=np.uint8)) |
| 60 | .long() |
| 61 | .to(self.device) |
| 62 | ) |
| 63 | sample = self.model.generate(inp[None, ...], length) |
| 64 | output_str = self.decode_tokens(sample[0]) |
| 65 | return output_str |
nothing calls this directly
no outgoing calls
no test coverage detected