| 9 | |
| 10 | |
| 11 | class CpmBeeLLM(LLM): |
| 12 | |
| 13 | model_name : str = "" |
| 14 | config: CPMBeeConfig = None |
| 15 | tokenizer: CPMBeeTokenizer = None |
| 16 | model: CPMBeeTorch = None |
| 17 | |
| 18 | def __init__(self, config_path: str, ckpt_path: str, device: str="cuda") -> None: |
| 19 | super().__init__() |
| 20 | self.model_name = ckpt_path |
| 21 | self.config = CPMBeeConfig.from_json_file(config_path) |
| 22 | self.tokenizer = CPMBeeTokenizer() |
| 23 | self.model = CPMBeeTorch(config=self.config) |
| 24 | |
| 25 | self.model.load_state_dict(torch.load(ckpt_path)) |
| 26 | if device == "cuda": |
| 27 | self.model.cuda() |
| 28 | |
| 29 | @property |
| 30 | def _llm_type(self) -> str: |
| 31 | return self.model_name |
| 32 | |
| 33 | def _call(self, prompt, stop: Optional[List[str]] = None) -> str: |
| 34 | # use beam search |
| 35 | beam_search = CPMBeeBeamSearch( |
| 36 | model=self.model, |
| 37 | tokenizer=self.tokenizer, |
| 38 | ) |
| 39 | inference_results = beam_search.generate([{"source":prompt, "<ans>":""}], max_length=512, repetition_penalty=1.2, beam_size=1) |
| 40 | output = inference_results[0]["<ans>"] |
| 41 | return output |
| 42 | |
| 43 | @property |
| 44 | def _identifying_params(self) -> Mapping[str, Any]: |
| 45 | """Get the identifying parameters.""" |
| 46 | return {"model_name": self.model_name} |
| 47 | |
| 48 | if __name__ == "__main__": |
| 49 | llm = CpmBeeLLM(config_path="path/to/cpm-bee/config.json", ckpt_path="path/to/cpm-bee/checkpoint/") |