diff --git a/modules/exllama.py b/modules/exllama.py index f3894b7a..f5cd2ae7 100644 --- a/modules/exllama.py +++ b/modules/exllama.py @@ -208,3 +208,8 @@ class ExllamaModel: ids = ids.view(1, -1) return self.tokenizer.decode(ids)[0] + + def get_logits(self, token_ids, **kwargs): + self.cache.current_seq_len = 0 + self.model.forward(token_ids[:, :-1], self.cache, input_mask=None, preprocess_only=True) + return self.model.forward(token_ids[:, -1:], self.cache, **kwargs).float().cpu() diff --git a/modules/exllamav2.py b/modules/exllamav2.py index 0bfe1f73..55903d80 100644 --- a/modules/exllamav2.py +++ b/modules/exllamav2.py @@ -113,3 +113,8 @@ class Exllamav2Model: ids = ids.view(1, -1) return self.tokenizer.decode(ids)[0] + + def get_logits(self, token_ids, **kwargs): + self.cache.current_seq_len = 0 + self.model.forward(token_ids[:, :-1], self.cache, input_mask=None, preprocess_only=True) + return self.model.forward(token_ids[:, -1:], self.cache, input_mask=None, **kwargs).float().cpu() diff --git a/modules/llamacpp_model.py b/modules/llamacpp_model.py index 5db6e27e..fc3a79f5 100644 --- a/modules/llamacpp_model.py +++ b/modules/llamacpp_model.py @@ -1,6 +1,7 @@ import re from functools import partial +import numpy as np import torch from modules import RoPE, shared @@ -100,6 +101,12 @@ class LlamaCppModel: def decode(self, tokens): return self.model.detokenize(tokens) + def get_logits(self, tokens): + self.model.eval(tokens) + logits = self.model._scores + logits = np.expand_dims(logits, 0) # batch dim is expected + return torch.tensor(logits, dtype=torch.float32) + def generate(self, prompt, state, callback=None): LogitsProcessorList = llama_cpp_lib().LogitsProcessorList diff --git a/modules/logits.py b/modules/logits.py index 3aed6624..d3b36a44 100644 --- a/modules/logits.py +++ b/modules/logits.py @@ -1,13 +1,32 @@ import torch from modules import sampler_hijack, shared +from modules.exllama import ExllamaModel +from modules.exllamav2 import Exllamav2Model +from modules.llamacpp_model import LlamaCppModel +from modules.logging_colors import logger from modules.text_generation import generate_reply global_scores = None def get_next_logits(prompt, state, use_samplers, previous): + if shared.model is None: + logger.error("No model is loaded! Select one in the Model tab.") + return 'Error: No model is loaded1 Select one in the Model tab.', previous + + is_non_hf_exllamav2 = isinstance(shared.model, Exllamav2Model) + is_non_hf_exllamav1 = isinstance(shared.model, ExllamaModel) + is_non_hf_llamacpp = isinstance(shared.model, LlamaCppModel) + if use_samplers: + if any([is_non_hf_exllamav2, is_non_hf_exllamav1, is_non_hf_llamacpp]): + logger.error("Sampler hijacking is not supported non-Huggingface loaders.") + # sampling is all done in c for exllama, so it is really hard to hijack + # it should be possible to hijack llamacpp sampler by hijacking all their sampling methods, + # but it is not implemented yet + return 'Error: Sampler hijacking is not supported non-Huggingface loaders. Please disable the "Use samplers" option.', previous + state['max_new_tokens'] = 1 state['auto_max_new_tokens'] = False for _ in generate_reply(prompt, state): @@ -15,17 +34,29 @@ def get_next_logits(prompt, state, use_samplers, previous): scores = sampler_hijack.global_scores[-1] else: - tokens = shared.tokenizer.encode(prompt, return_tensors='pt').cuda() - output = shared.model(input_ids=tokens) - scores = output['logits'][-1][-1] + if is_non_hf_exllamav2 or is_non_hf_exllamav1: + tokens = shared.tokenizer.encode(prompt).cuda() + scores = shared.model.get_logits(tokens)[-1][-1] + elif is_non_hf_llamacpp: + tokens = shared.tokenizer.encode(prompt) + scores = shared.model.get_logits(tokens)[-1][-1] + else: + tokens = shared.tokenizer.encode(prompt, return_tensors='pt').cuda() + output = shared.model(input_ids=tokens) + scores = output['logits'][-1][-1] probs = torch.softmax(scores, dim=-1, dtype=torch.float) topk_values, topk_indices = torch.topk(probs, k=25, largest=True, sorted=True) topk_values = [f"{float(i):.5f}" for i in topk_values] + if is_non_hf_exllamav1 or is_non_hf_llamacpp: + topk_indices = [i.expand((1, 1)) for i in topk_indices] + tokens = [shared.tokenizer.decode(i) for i in topk_indices] + if is_non_hf_llamacpp: + tokens = [i.decode('utf-8') for i in tokens] # llamacpp returns bytes, not str output = '' for row in list(zip(topk_values, tokens)): - output += f"{row[0]} - {row[1]}\n" + output += f"{row[0]} - {repr(row[1])[1:-1]}\n" return output, previous diff --git a/modules/text_generation.py b/modules/text_generation.py index a3755c10..41f12ddc 100644 --- a/modules/text_generation.py +++ b/modules/text_generation.py @@ -150,7 +150,7 @@ def get_token_ids(prompt): output = '' for row in list(zip(tokens, decoded_tokens)): - output += f"{str(int(row[0])).ljust(5)} - {row[1]}\n" + output += f"{str(int(row[0])).ljust(5)} - {repr(row[1])[1:-1]}\n" return output