import gradio as gr
from transformers import pipeline
import ast

pipe = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", device=-1)

createListSysPrompt = (
    "Please create a Python dictionary named 'list' that you will use to keep track of whatever is asked for. "
    "The format should be exactly like this:\n\n"
    "list = {\n"
    "    'typeOfThing1': ['exOfThing1_1', 'exOfThing1_2'],\n"
    "    'typeOfThing2': ['exOfThing2_1'],\n"
    "}\n\n"
    "Always return a valid Python dictionary with keys as categories and values as lists of items."
)

updateListSysPrompt = (
   "If you see mention of a thing of a type in the list that's not in the list, add it to the list."
)

def updateList(newText):
  preparePrompt = []
  list = []

  if list == []:
    preparePrompt.append({"role": "system", "content": createListSysPrompt})
  else:
    preparePrompt.append({"role": "system", "content": updateListSysPrompt})

  preparePrompt.append({"role": "user", "content": newText})

  prompt = pipe.tokenizer.apply_chat_template(preparePrompt, tokenize=False, add_generation_prompt=True)
  outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
  return outputs[0]["generated_text"]


demo = gr.Interface(fn=updateList, inputs="textbox", outputs="textbox")

if __name__ == "__main__":
    demo.launch()