switch to pipelines

app.py

@@ -1,28 +1,11 @@
 # https://www.gradio.app/guides/using-hugging-face-integrations
 
-from transformers import pipeline
 import gradio as gr
-
-model = "TinyLlama/TinyLlama-1.1B-Chat-v0.3"
-
-pipe = pipeline("conversational", model=model)
-demo = gr.Interface.from_pipeline(pipe)
-demo.launch()
-
-'''
-demo = gr.Interface.load(model)
-'''
-
-
-"""
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import gradio as gr
-import torch
+from transformers import pipeline, Conversation
 
 model = "mistralai/Mistral-7B-Instruct-v0.1"
 model = "TinyLlama/TinyLlama-1.1B-Chat-v0.3"
 
-# Gradio
 title = "Shisa 7B"
 description = "Test out Shisa 7B in either English or Japanese."
 placeholder = "Type Here / ここに入力してください" 
@@ -33,42 +16,28 @@ examples = [
     "こんにちは、いかがお過ごしですか？",
 ]
 
-tokenizer = AutoTokenizer.from_pretrained(model)
-model = AutoModelForCausalLM.from_pretrained(model)
+# Docs: https://github.com/huggingface/transformers/blob/main/src/transformers/pipelines/conversational.py
+conversation = Conversation()
+chatbot = pipeline('conversational', model)
+'''
+conversation = Conversation("Going to the movies tonight - any suggestions?")
+conversation.add_message({"role": "assistant", "content": "The Big lebowski."})
+conversation.add_message({"role": "user", "content": "Is it good?"})
+conversation.messages[:-1]
+'''
 
 def chat(input, history=[]):
-    input_ids = tokenizer.encode(input + tokenizer.eos_token, return_tensors="pt")
-    history = model.generate(
-        input_ids, max_length=4000, pad_token_id=tokenizer.eos_token_id
-    ).tolist()
-
-    # convert the tokens to text, and then split the responses into lines
-    response = tokenizer.decode(history[0]).split("<|endoftext|>")
-    '''
-    # tokenize the new input sentence
-    new_user_input_ids = tokenizer.encode(
-        input + tokenizer.eos_token, return_tensors="pt"
-    )
-
-    # append the new user input tokens to the chat history
-    bot_input_ids = torch.cat([torch.LongTensor(history), new_user_input_ids], dim=-1)
-
-    # generate a response
-    history = model.generate(
-        bot_input_ids, max_length=4000, pad_token_id=tokenizer.eos_token_id
-    ).tolist()
-
-    # convert the tokens to text, and then split the responses into lines
-    response = tokenizer.decode(history[0]).split("<|endoftext|>")
-    # print('decoded_response-->>'+str(response))
-    response = [
-        (response[i], response[i + 1]) for i in range(0, len(response) - 1, 2)
-    ]  # convert to tuples of list
-    # print('response-->>'+str(response))
-    '''
+    conversation.add_message({"role": "user", "content": input})
+    # we do this shuffle so local shadow response doesn't get created
+    response_conversation = chatbot(conversation)
+    print(response_conversation)
+    print(response_conversation.messages)
+    print(response_conversation.messages[-1]["content"])
+
+    conversation.add_message(response_conversation.messages[-1])
+    response = conversation.messages[-1]["content"]
     return response, history
 
-
 gr.ChatInterface(
     chat,
     chatbot=gr.Chatbot(height=400),
@@ -80,5 +49,31 @@ gr.ChatInterface(
     cache_examples=False,
     undo_btn="Delete Previous",
     clear_btn="Clear",
-).queue().launch()
-"""
+).launch()
+
+'''
+gr.Interface.load(
+    "EleutherAI/gpt-j-6B",
+    inputs=gr.Textbox(lines=5, label="Input Text"),
+    title=title,
+    description=description,
+    article=article,
+).launch()
+
+
+# Doesn't support conversational pipelin
+pipe = pipeline('conversational', model)
+gr.Interface.from_pipeline(pipe).launch()
+
+
+'''
+
+
+# For async
+# ).queue().launch()
+
+'''
+# Pipeline doesn't support conversational...
+pipe = pipeline("conversational", model=model)
+demo = gr.Interface.from_pipeline(pipe)
+'''