Spaces:
Sleeping
Sleeping
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, StoppingCriteria, StoppingCriteriaList | |
| import time | |
| import numpy as np | |
| from torch.nn import functional as F | |
| import os | |
| from .base_model import BaseLLMModel | |
| class StopOnTokens(StoppingCriteria): | |
| def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool: | |
| stop_ids = [50278, 50279, 50277, 1, 0] | |
| for stop_id in stop_ids: | |
| if input_ids[0][-1] == stop_id: | |
| return True | |
| return False | |
| class StableLM_Client(BaseLLMModel): | |
| def __init__(self, model_name) -> None: | |
| super().__init__(model_name=model_name) | |
| print(f"Starting to load StableLM to memory") | |
| self.model = AutoModelForCausalLM.from_pretrained( | |
| "stabilityai/stablelm-tuned-alpha-7b", torch_dtype=torch.float16).cuda() | |
| self.tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-tuned-alpha-7b") | |
| self.generator = pipeline('text-generation', model=self.model, tokenizer=self.tokenizer, device=0) | |
| print(f"Sucessfully loaded StableLM to the memory") | |
| self.system_prompt = """StableAssistant | |
| - StableAssistant is A helpful and harmless Open Source AI Language Model developed by Stability and CarperAI. | |
| - StableAssistant is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user. | |
| - StableAssistant is more than just an information source, StableAssistant is also able to write poetry, short stories, and make jokes. | |
| - StableAssistant will refuse to participate in anything that could harm a human.""" | |
| def user(self, user_message, history): | |
| history = history + [[user_message, ""]] | |
| return "", history, history | |
| def bot(self, history, curr_system_message): | |
| messages = f"<|SYSTEM|># {self.system_prompt}" + \ | |
| "".join(["".join(["<|USER|>"+item[0], "<|ASSISTANT|>"+item[1]]) | |
| for item in history]) | |
| output = self.generate(messages) | |
| history[-1][1] = output | |
| time.sleep(1) | |
| return history, history | |
| def _get_stablelm_style_input(self): | |
| messages = self.system_prompt + \ | |
| "".join(["".join(["<|USER|>"+self.history[i]["content"], "<|ASSISTANT|>"+self.history[i + 1]["content"]]) | |
| for i in range(0, len(self.history), 2)]) | |
| return messages | |
| def generate(self, text, bad_text=None): | |
| stop = StopOnTokens() | |
| result = self.generator(text, max_new_tokens=1024, num_return_sequences=1, num_beams=1, do_sample=True, | |
| temperature=1.0, top_p=0.95, top_k=1000, stopping_criteria=StoppingCriteriaList([stop])) | |
| return result[0]["generated_text"].replace(text, "") | |
| def contrastive_generate(self, text, bad_text): | |
| with torch.no_grad(): | |
| tokens = self.tokenizer(text, return_tensors="pt")[ | |
| 'input_ids'].cuda()[:, :4096-1024] | |
| bad_tokens = self.tokenizer(bad_text, return_tensors="pt")[ | |
| 'input_ids'].cuda()[:, :4096-1024] | |
| history = None | |
| bad_history = None | |
| curr_output = list() | |
| for i in range(1024): | |
| out = self.model(tokens, past_key_values=history, use_cache=True) | |
| logits = out.logits | |
| history = out.past_key_values | |
| bad_out = self.model(bad_tokens, past_key_values=bad_history, | |
| use_cache=True) | |
| bad_logits = bad_out.logits | |
| bad_history = bad_out.past_key_values | |
| probs = F.softmax(logits.float(), dim=-1)[0][-1].cpu() | |
| bad_probs = F.softmax(bad_logits.float(), dim=-1)[0][-1].cpu() | |
| logits = torch.log(probs) | |
| bad_logits = torch.log(bad_probs) | |
| logits[probs > 0.1] = logits[probs > 0.1] - bad_logits[probs > 0.1] | |
| probs = F.softmax(logits) | |
| out = int(torch.multinomial(probs, 1)) | |
| if out in [50278, 50279, 50277, 1, 0]: | |
| break | |
| else: | |
| curr_output.append(out) | |
| out = np.array([out]) | |
| tokens = torch.from_numpy(np.array([out])).to( | |
| tokens.device) | |
| bad_tokens = torch.from_numpy(np.array([out])).to( | |
| tokens.device) | |
| return self.tokenizer.decode(curr_output) | |
| def get_answer_at_once(self): | |
| messages = self._get_stablelm_style_input() | |
| return self.generate(messages) | |