import gradio as gr
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
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)