Spaces:
Sleeping
Sleeping
| from tokenize import tokenize | |
| import torch | |
| import torch.nn.functional as F | |
| from transformers import BertTokenizer, GPT2LMHeadModel | |
| def replace_special_tokens(tokens:list, special_token_dict:dict)->list: | |
| replaced_tokens:list = [] | |
| for token in tokens: | |
| if token in special_token_dict.keys(): | |
| replaced_tokens.append( special_token_dict[token] ) | |
| else: | |
| replaced_tokens.append( token ) | |
| return replaced_tokens | |
| def top_k_top_p_filtering(logits, top_k:int=0, top_p:float=0.0, filter_value:float=-float('Inf')): | |
| top_k = min( top_k, logits.size(-1) ) | |
| if top_k > 0: | |
| indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None] | |
| logits[indices_to_remove] = filter_value | |
| if top_p > 0.0: | |
| sorted_logits, sorted_indices = torch.sort(logits, descending=True) | |
| cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) | |
| sorted_indices_to_remove = cumulative_probs > top_p | |
| sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() | |
| sorted_indices_to_remove[..., 0] = 0 | |
| indices_to_remove = sorted_indices[sorted_indices_to_remove] | |
| logits[indices_to_remove] = filter_value | |
| return logits | |
| class ChatBot(): | |
| def get_chat_bot(checkpoint:str, vocab_path:str = None, special_token_list:list = [])->object: | |
| tokenizer = ChatBot.get_tokenizer(checkpoint, vocab_path, special_token_list) | |
| model = GPT2LMHeadModel.from_pretrained(checkpoint) | |
| return ChatBot(tokenizer, model) | |
| def get_tokenizer(checkpoint:str, vocab_path:str = None, special_token_list:list = [])->object: | |
| if vocab_path is None: | |
| tokenizer = BertTokenizer.from_pretrained(checkpoint, sep_token="[SEP]", pad_token="[PAD]", cls_token="[CLS]") | |
| else: | |
| tokenizer = BertTokenizer(vocab_path, sep_token="[SEP]", pad_token="[PAD]", cls_token="[CLS]") | |
| tokenizer.add_special_tokens( {'additional_special_tokens':special_token_list} ) | |
| return tokenizer | |
| def __init__(self, tokenizer:object, model:object)->None: | |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
| self.tokenizer = tokenizer | |
| self.model = model | |
| self.model.to(self.device) | |
| def convert_ids_to_tokens(self, ids:list): | |
| return self.tokenizer.convert_ids_to_tokens(ids) | |
| def convert_ids_to_text(self, ids): | |
| return "".join( self.convert_ids_to_tokens(ids) ) | |
| def convert_text_to_ids(self, text, add_special_tokens=False): | |
| return self.tokenizer.encode(text, add_special_tokens=add_special_tokens) | |
| def get_prediction(self, input_tensor, input_ids, repetition_penalty, temperature, top_k, top_p): | |
| self.model.eval() | |
| generated_ids = [] | |
| for _ in range(64): | |
| output_pt = self.model(input_tensor) | |
| next_token_logits = output_pt.logits[0, -1, :] | |
| for id in set(input_ids): | |
| if id != self.tokenizer.sep_token_id: | |
| next_token_logits[id] /= repetition_penalty | |
| for id in set(generated_ids): | |
| next_token_logits[id] /= repetition_penalty | |
| next_token_logits = next_token_logits / temperature | |
| next_token_logits[self.tokenizer.convert_tokens_to_ids('[UNK]')] = -float('Inf') | |
| filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=top_k, top_p=top_p) | |
| next_token = torch.multinomial( F.softmax(filtered_logits, dim=-1), num_samples=1 ) | |
| if next_token == self.tokenizer.sep_token_id: | |
| break | |
| input_tensor = torch.cat( (input_tensor, next_token.unsqueeze(0)), dim=1 ) | |
| generated_ids.append( next_token.item() ) | |
| return generated_ids | |
| def chat(self:object, text:str, history:list, role_card:dict={}) -> str: | |
| text_ids = self.tokenizer.encode(text, add_special_tokens=False) | |
| history.append(text_ids) | |
| input_ids = [self.tokenizer.cls_token_id] | |
| for history_utr in history[-50:]: | |
| input_ids.extend(history_utr) | |
| input_ids.append(self.tokenizer.sep_token_id) | |
| input_tensor = torch.tensor(input_ids).to(self.device).unsqueeze(0) | |
| generated_ids = self.get_prediction(input_tensor, input_ids, repetition_penalty=1.2, temperature=0.73, top_k=10, top_p=0.7) | |
| history.append(generated_ids) | |
| generated_tokens = replace_special_tokens( self.convert_ids_to_tokens(generated_ids), role_card ) | |
| return "".join(generated_tokens), history | |