Update app.py

app.py

@@ -7,103 +7,103 @@ 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)
+# 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)