import chatglm_cpp
import gradio as gr
from pathlib import Path

model_file_path = "chatglm3-ggml_q4_0.bin"
chatglm_llm = chatglm_cpp.Pipeline(Path(model_file_path))

def predict(message, history):
    '''
    messages = []
    for human_content, system_content in history:
        message_human = {
            "role": "user",
            "content": human_content + "\n",
        }
        message_system = {
            "role": "system",
            "content": system_content + "\n",
        }
        messages.append(message_human)
        messages.append(message_system)
    message_human = {
        "role": "user",
        "content": message + "\n",
    }
    messages.append(message_human)
    # Llamaでの回答を取得（ストリーミングオン）
    streamer = llama.create_chat_completion(messages, stream=True)
    '''
    flatten_history = []
    for a, b in history:
        flatten_history.append(a)
        flatten_history.append(b)

    streamer = chatglm_llm.chat(
        history= flatten_history + [message], do_sample=False,
        stream = True
        )

    '''
    partial_message = ""
    for msg in streamer:
        message = msg['choices'][0]['delta']
        if 'content' in message:
            partial_message += message['content']
            yield partial_message
    '''
    response = ""
    for new_text in streamer:
        response += new_text
        yield response

gr.ChatInterface(predict,
    chatbot=gr.Chatbot(height=300),
    textbox=gr.Textbox(placeholder="你好 人工智能助手 ChatGLM3，我可以问你一些问题吗？", container=False, scale=7),
    title="ChatGLM3 Chatbot 🐼",
    description="与人工智能助手 ChatGLM3 进行对话",
    theme="soft",
    examples=[
            "哈利波特和赫敏是什么关系？",
        "请解释下面的emoji符号描述的情景👨👩🔥❄️",
            "明朝内阁制度的特点是什么?",
            "如何进行经济建设?", 
             "你听说过马克思吗？",
        ],
    cache_examples=False,
    retry_btn=None,
    undo_btn="Delete Previous",
    clear_btn="Clear",
).launch("0.0.0.0")