import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
modelname="EleutherAI/gpt-neo-2.7B"
config = AutoConfig.from_pretrained(modelname)
tokenizer = AutoTokenizer.from_pretrained(modelname)
model = AutoModelForCausalLM.from_pretrained(modelname,config=config).to("cuda")


def botsay(user_input):
  prompt = "This is a conversation between Human and AI bot. AI's name is ThatGPT."  
  new_token_id=None
  gen_tokens=""
  new_token=""
  j =6
  length=0
  limit = 128
  thatid=5562
  cont = True
  last_apppended = False
  cnt=0
  disable_repeat_length= 5
  disable_repeat_count = 2
  tokens=[]
  while(cont):
    cnt+=1
    prob = 1.0
    input_ids=tokenizer(prompt+user_input+"\nAI:"+gen_tokens,return_tensors="pt").input_ids
    length=len(input_ids)
    if length >limit:
      gen_tokens="⚠️sorry length limit. please reload the browser."
      return gen_tokens
    outs=model(input_ids=input_ids.to("cuda"))
    topk = torch.topk(outs.logits.squeeze()[-1,:],k=j+1).indices
    if new_token =="that":
      that_id = 326
    elif new_token ==" that":
      that_id = -1
    elif new_token[-1:] ==" ":
      that_id = 5562
    else:
      that_id = 326

    if ("thatGPT" in gen_tokens[-12:]):
      that_id = -1
    if last_apppended:
      that_id = -1
    if that_id in topk:
      new_token_id = that_id
    else:
      new_token_id = torch.argmax(outs.logits.squeeze()[-1,:])
      new_token=tokenizer.decode(new_token_id)
    new_token=tokenizer.decode(new_token_id)
    prev_tokens=gen_tokens
    gen_tokens+=new_token
    if (cnt>10) and (disable_repeat_count<gen_tokens.count(gen_tokens[-disable_repeat_length:])):
      gen_tokens=prev_tokens
      new_token = tokenizer.decode(topk[torch.randint(5, (1,1)).item()])
      gen_tokens+=new_token

    if new_token_id==50256 or new_token_id==198 or new_token=="<|endoftext|>":
      if ("that" not in gen_tokens):
        gen_tokens = gen_tokens.replace("\n","").replace(".","")
        gen_tokens += " that"
      else:
        cont = False
  return gen_tokens.replace("<br>","").replace("AI:","").replace("\xa0","")


    

import gradio as gr
def add_text(history, text):
    history = history + [(text, None)]
    return history, ""


def bot(history):
    serial_history=""
    for h in history:
      serial_history+="\nHuman:"+h[0]
      if h[1]==None:
        break
      serial_history+="\nAI:"+h[1].replace("<br>","")

    response = botsay(serial_history)
    history[-1][1] = response
    serial_history+="\nAI:"+response
    return history

with gr.Blocks() as demo:
    gr.Markdown("# ThatGPT")
    chatbot = gr.Chatbot([], elem_id="chatbot").style(height=750)

    with gr.Row():
        with gr.Column(scale=0.85):
            txt = gr.Textbox(
                show_label=False,
                placeholder="AI always replies with \"that\".",
            ).style(container=False)
        
    txt.submit(add_text, [chatbot, txt], [chatbot, txt]).then(
        bot, chatbot, chatbot
    )

demo.launch()