import gradio as gr
import random

text_to_text_gm130b = gr.Blocks.load(name="spaces/THUDM/GLM-130B")

def block_inference(prompt):
  #thudm glm 130 b space 
  #generated_text = text_to_text_gm130b(model_input=prompt, seed=1234, out_seq_length=256, min_gen_length=0, sampling_strategy='BeamSearchStrategy', num_beams=2, length_penalty=1, no_repeat_ngram_size=3, temperature=0.7, topk=1, topp=0)
  generated_text = text_to_text_gm130b(prompt, 1234, 256, 0, 'BeamSearchStrategy', 2, 1, 3, 0.7, 1, 0)
  return generated_text
  
def chat(message, history):
  history = history or []
  message = message.lower()
  response = block_inference(message)
  #if message.startswith("how many"):
  #    response = random.randint(1, 10)
  #elif message.startswith("how"):
  #    response = random.choice(["Great", "Good", "Okay", "Bad"])
  #elif message.startswith("where"):
  #    response = random.choice(["Here", "There", "Somewhere"])
  #else:
  #    response = "I don't know"
  history.append((message, response))
  return history, history

chatbot = gr.Chatbot().style(color_map=("green", "pink"))
demo = gr.Interface(
    chat,
    ["text", "state"],
    [chatbot, "state"],
    allow_flagging="never",
)
if __name__ == "__main__":
    demo.launch()