perf: assign to candidates data structure instead

This commit is contained in:
Andrei Betlen 2023-07-07 18:58:43 -04:00
parent ca11673061
commit a14d8a9b3f

View file

@ -487,9 +487,9 @@ class Llama:
nl_logit = logits[self._token_nl] nl_logit = logits[self._token_nl]
candidates = self._candidates candidates = self._candidates
candidates_data = self._candidates_data candidates_data = self._candidates_data
candidates_data["id"] = np.arange(n_vocab, dtype=np.intc) # type: ignore candidates_data["id"][:] = np.arange(n_vocab, dtype=np.intc) # type: ignore
candidates_data["logit"] = logits candidates_data["logit"][:] = logits
candidates_data["p"] = np.zeros(n_vocab, dtype=np.single) candidates_data["p"][:] = np.zeros(n_vocab, dtype=np.single)
candidates.data = candidates_data.ctypes.data_as(llama_cpp.llama_token_data_p) candidates.data = candidates_data.ctypes.data_as(llama_cpp.llama_token_data_p)
candidates.sorted = llama_cpp.c_bool(False) candidates.sorted = llama_cpp.c_bool(False)
candidates.size = llama_cpp.c_size_t(n_vocab) candidates.size = llama_cpp.c_size_t(n_vocab)