Spaces:
Sleeping
Sleeping
| 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() | |