Adding app.py

app.py

@@ -4,106 +4,100 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, Stopping
 import time
 import numpy as np
 from torch.nn import functional as F
-import os
-token_key = os.environ.get("HF_ACCESS_TOKEN")
-
-# if torch.cuda.is_available():
-# m = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-tuned-alpha-7b",use_auth_token=token_key, torch_dtype=torch.float16).cuda()
-# tok = AutoTokenizer.from_pretrained("stabilityai/stablelm-tuned-alpha-7b",use_auth_token=token_key)
-# else:
-# m = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-tuned-alpha-7b",use_auth_token=token_key, torch_dtype=torch.float16)
-# tok = AutoTokenizer.from_pretrained("stabilityai/stablelm-tuned-alpha-7b",use_auth_token=token_key)
-# generator = pipeline('text-generation', model=m, tokenizer=tok, device=0)
-
-
-# start_message = """<|SYSTEM|># 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."""
-
-
-# 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
-
-
-# def contrastive_generate(text, bad_text):
-# with torch.no_grad():
-# if torch.cuda_is_available():
-# tokens = tok(text, return_tensors="pt")['input_ids'].cuda()[:,:4096-1024]
-# bad_tokens = tok(bad_text, return_tensors="pt")['input_ids'].cuda()[:,:4096-1024]
-# else:
-# tokens = tok(text, return_tensors="pt")['input_ids'][:,:4096-1024]
-# bad_tokens = tok(bad_text, return_tensors="pt")['input_ids'][:,:4096-1024]
-# history = None
-# bad_history = None
-# curr_output = list()
-# for i in range(1024):
-# out = m(tokens, past_key_values=history, use_cache=True)
-# logits = out.logits
-# history = out.past_key_values
-# bad_out = m(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 tok.decode(curr_output)
-
-# def generate(text, bad_text=None):
-# stop = StopOnTokens()
-# result = 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 user(user_message, history):
-# return "", history + [[user_message, ""]]
-
-
-# def bot(history, curr_system_message):
-# messages = curr_system_message + "".join(["".join(["<|USER|>"+item[0], "<|ASSISTANT|>"+item[1]]) for item in history])
-# output = generate(messages)
-# history[-1][1] = output
-# time.sleep(1)
-# return history
-
-
-# def system_update(msg):
-# global curr_system_message
-# curr_system_message = msg
-
-
-# with gr.Blocks() as demo:
-# gr.Markdown("###StableLM-tuned-Alpha-7B Chat")
-# with gr.Row():
-# with gr.Column():
-# chatbot = gr.Chatbot([])
-# clear = gr.Button("Clear")
-# with gr.Column():
-# system_msg = start_message#gr.Textbox(start_message, label="System Message", interactive=True)
-# msg = gr.Textbox(label="Chat Message")
-
-# msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-# bot, [chatbot, system_msg], chatbot
-# )
-# system_msg.change(system_update, system_msg, None, queue=False)
-# clear.click(lambda: None, None, chatbot, queue=False)
-# demo.launch(share=True)
+import os 
+auth_key = os.environ["HF_ACCESS_TOKEN"]
+print(f"Starting to load the model to memory")
+m = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-tuned-alpha-7b",use_auth_token=auth_key, torch_dtype=torch.float16).cuda()
+tok = AutoTokenizer.from_pretrained("stabilityai/stablelm-tuned-alpha-7b",use_auth_token=auth_key)
+generator = pipeline('text-generation', model=m, tokenizer=tok, device=0)
+print(f"Sucessfully loaded the model to the memory")
+
+start_message = """<|SYSTEM|># 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."""
+
+
+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
+
+
+def contrastive_generate(text, bad_text):
+ with torch.no_grad():
+ tokens = tok(text, return_tensors="pt")['input_ids'].cuda()[:,:4096-1024]
+ bad_tokens = tok(bad_text, return_tensors="pt")['input_ids'].cuda()[:,:4096-1024]
+ history = None
+ bad_history = None
+ curr_output = list()
+ for i in range(1024):
+ out = m(tokens, past_key_values=history, use_cache=True)
+ logits = out.logits
+ history = out.past_key_values
+ bad_out = m(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 tok.decode(curr_output)
+
+def generate(text, bad_text=None):
+ stop = StopOnTokens()
+ result = 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 user(user_message, history):
+ return "", history + [[user_message, ""]]
+
+
+def bot(history, curr_system_message):
+ messages = curr_system_message + "".join(["".join(["<|USER|>"+item[0], "<|ASSISTANT|>"+item[1]]) for item in history])
+ output = generate(messages)
+ history[-1][1] = output
+ time.sleep(1)
+ return history
+
+
+def system_update(msg):
+ global curr_system_message
+ curr_system_message = msg
+
+updated_system_message = ""
+
+with gr.Blocks() as demo:
+ gr.Markdown("## StableLM-Tuned-Alpha-7b Chat")
+ with gr.Row():
+ chatbot = gr.Chatbot([])
+ clear = gr.Button("Clear Chat History")
+ system_msg = gr.Textbox(start_message, label="System Message", interactive=False,visible=False)
+ #system_msg = start_message
+ msg = gr.Textbox(label="Chat Message Box")
+
+ msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+ bot, [chatbot, system_msg], chatbot
+ )
+ system_update(system_msg)
+ system_msg.change(system_update, system_msg, None, queue=False)
+ clear.click(lambda: None, None, chatbot, queue=False)
+
+demo.launch()