import json 
import gradio as gr
from llama_cpp import Llama



system_message = "\nBelow is an instruction that describes a task. Write a response that appropriately completes the request."

title = "MY First Llama2 7B Chatbot"
description = """
基于6000条问答数据微调过后的Llama2，初步激活了Llama2的中文能力，使其能够说中文。
"""
css = """.toast-wrap { display: none !important } """
examples=[
    '使用python写一个post请求',
    '使用swiftui自定义图片大小',
    ]
# 
llm = Llama(model_path="ggml-model-q4_0.bin", n_ctx=2048)

def predict(message, chatbot):
    
    input_prompt = f"[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n "
    for interaction in chatbot:
        input_prompt = input_prompt + str(interaction[1]) + " [/INST] " + str(interaction[0]) + " </s><s> [INST] "
    "Question: What are the names of the planets in the solar system? Answer: ",
    input_prompt = input_prompt + str(message) + " [/INST] "
    stream = llm(
    input_prompt, 
    max_tokens=512,
    stop=["[INST]"],
    stream=True)

    partial_message = ""
    for line in stream:
        # Load as JSON
        try:
            if 'choices' in line:
                partial_message = partial_message + line["choices"][0]["text"]
                yield partial_message
            else:
                gr.Warning(f"The key 'token' does not exist in this JSON object: {line}")

        except json.JSONDecodeError:
            gr.Warning(f"This line is not valid JSON: {line}")
            continue
        except KeyError as e:
            gr.Warning(f"KeyError: {e} occurred for JSON object: {line}")
            continue





# Gradio Demo 
with gr.Blocks() as demo:

    with gr.Tab("Streaming"):
        gr.ChatInterface(predict, title=title, description=description, css=css, examples=examples, cache_examples=False) 
        # gr.ChatInterface(predict, title=title, css=css, examples=examples, cache_examples=False) 
demo.queue(concurrency_count=75).launch(debug=True)