import gradio as gr
from llama_cpp import Llama
import os
import json

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
## Imports
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
# Initialize the Llama model

## Download the GGUF model
model_name = "johnpaulbin/qwen2.5-3b-tokipona-v4-merged-fp16-chat-Q8_0-GGUF"
model_file = "qwen2.5-3b-tokipona-v4-merged-fp16-chat-q8_0.gguf"  # This is the specific model file we'll use in this example.
model_path = hf_hub_download(model_name, filename=model_file)

## Instantiate model from downloaded file
llm = Llama(
    model_path=model_path,
    n_ctx=1024,           # Context length to use
    n_threads=2,          # Number of CPU threads to use
    n_gpu_layers=-1,      # Number of model layers to offload to GPU
    chat_format="chatml",
    flash_attn=True
)

# Initialize the conversation with a system prompt
system_prompt = "sina ilo Elelen. sina jan pona pi wile pali. sina pona tawa jan ante li wile pali pona tawa ona."
initial_messages = [{"role": "system", "content": system_prompt}]

# Initialize the dataset
dataset = {
    "prompt": [],
    "completion": [],
    "label": []
}

def generate_response(user_message, chat_history, messages):
    # Append the user's message to the conversation history
    messages = messages + [{"role": "user", "content": user_message}]
    print(messages)
    # Call the model to get assistant's response
    try:
        response = llm.create_chat_completion(
            messages=messages,
            temperature=0.0,
            top_p=1,
            top_k=1,
            frequency_penalty=1
        )
        assistant_response = response['choices'][0]['message']['content']
        # Append the assistant's response to the conversation history
        messages = messages + [{"role": "assistant", "content": assistant_response}]

        # Update the chat history
        chat_history = chat_history + [(user_message, assistant_response)]
        return "", chat_history, messages
    except ValueError as e:
        print(f"Error: You've reached the max context length for the model. Clear the conversation to continue typing.")

        assistant_response = "Error: You've reached the max context length for the model. Clear the conversation to continue typing."
        # Append the assistant's response to the conversation history
        messages = messages + [{"role": "assistant", "content": assistant_response}]

        # Update the chat history
        chat_history = chat_history + [(user_message, assistant_response)]
        return "", chat_history, messages

def upvote_click(chat_history, messages, dataset):
    notification = "Your feedback has been submitted! Feel free to keep chatting or clear the history."
    if len(messages) >= 2:
        # Construct the full conversation history
        full_conversation = ""
        for msg in messages:
            if msg['role'] == 'system':
                full_conversation += f"
system\n{msg['content']}
\n"
            elif msg['role'] == 'user':
                full_conversation += f"
user\n{msg['content']}
\n"
            elif msg['role'] == 'assistant':
                full_conversation += f"
assistant\n{msg['content']}
"

        # Split the conversation into prompt and completion
        prompt_parts = full_conversation.rsplit("
assistant", 1)
        prompt = prompt_parts[0]
        completion = f"
assistant{prompt_parts[1]}" if len(prompt_parts) > 1 else ""

        # Read existing data from feedback.json
        try:
            with open('feedback.json', 'r') as f:
                existing_data = json.load(f)
        except FileNotFoundError:
            existing_data = {'prompt': [], 'completion': [], 'label': []}

        # Check for duplicates
        duplicate = False
        for p, c in zip(existing_data['prompt'], existing_data['completion']):
            if p == prompt and c == completion:
                duplicate = True
                break

        if not duplicate:
            existing_data['prompt'].append(prompt)
            existing_data['completion'].append(completion)
            existing_data['label'].append(True)
            # Save to feedback.json
            with open('feedback.json', 'w') as f:
                json.dump(existing_data, f)
            dataset = existing_data
            print("Upvoted the response! YAAY!!!")
        else:
            notification = "You have already submitted feedback for this conversation."
    else:
        notification = "No conversation to submit feedback for."
    return dataset, notification

def downvote_click(chat_history, messages, dataset):
    notification = "Your feedback has been submitted! Feel free to keep chatting or clear the history."
    if len(messages) >= 2:
        # Construct the full conversation history
        full_conversation = ""
        for msg in messages:
            if msg['role'] == 'system':
                full_conversation += f"
system\n{msg['content']}
\n"
            elif msg['role'] == 'user':
                full_conversation += f"
user\n{msg['content']}
\n"
            elif msg['role'] == 'assistant':
                full_conversation += f"
assistant\n{msg['content']}
"

        # Split the conversation into prompt and completion
        prompt_parts = full_conversation.rsplit("
assistant", 1)
        prompt = prompt_parts[0]
        completion = f"
assistant{prompt_parts[1]}" if len(prompt_parts) > 1 else ""

        # Read existing data from feedback.json
        try:
            with open('feedback.json', 'r') as f:
                existing_data = json.load(f)
        except FileNotFoundError:
            existing_data = {'prompt': [], 'completion': [], 'label': []}

        # Check for duplicates
        duplicate = False
        for p, c in zip(existing_data['prompt'], existing_data['completion']):
            if p == prompt and c == completion:
                duplicate = True
                break

        if not duplicate:
            existing_data['prompt'].append(prompt)
            existing_data['completion'].append(completion)
            existing_data['label'].append(False)
            # Save to feedback.json
            with open('feedback.json', 'w') as f:
                json.dump(existing_data, f)
            dataset = existing_data
            print("Downvoted the response!")
        else:
            notification = "You have already submitted feedback for this conversation."
    else:
        notification = "No conversation to submit feedback for."
    return dataset, notification

def clear_history():
    return [], initial_messages

def retry_response(chat_history, messages):
    notification = ""
    # Check if there is an assistant response to retry
    if len(messages) >= 2 and messages[-1]['role'] == 'assistant':
        # Remove the last assistant message from messages
        messages = messages[:-1]
        # Remove the last exchange from chat_history
        if len(chat_history) >= 1:
            chat_history = chat_history[:-1]
        # Regenerate the assistant's response
    try:
        response = llm.create_chat_completion(
            messages=messages,
            temperature=0.2,
            top_p=0.95,
            top_k=5,
            frequency_penalty=1
        )
        assistant_response = response['choices'][0]['message']['content']
        # Append the assistant's response to the conversation history
        messages = messages + [{"role": "assistant", "content": assistant_response}]
        chat_history = chat_history + [(messages[-2]['content'], assistant_response)]
        return chat_history, messages, notification
    except ValueError as e:
        assistant_response = "Error: You've reached the max context length for the model. Clear the conversation to continue typing."
        # Append the assistant's response to the conversation history
        messages = messages + [{"role": "assistant", "content": assistant_response}]

        # Update the chat history
        chat_history = chat_history + [(user_message, assistant_response)]
        return chat_history, messages, notification
    else:
        notification = "Nothing to retry."
    return chat_history, messages, notification

with gr.Blocks() as demo:
    gr.Markdown("# Toki Pona LLM")
    gr.Markdown("Help improve the model by voting after testing it! Note: The model best performs in a toki-pona-only conversation.")
    messages_state = gr.State(initial_messages)
    dataset_state = gr.State(dataset)
    chat_history = gr.Chatbot()
    with gr.Row():
        user_message = gr.Textbox(placeholder="Enter your message")
        send_button = gr.Button("Send")
    with gr.Row():
        upvote_button = gr.Button("Upvote 👍")
        downvote_button = gr.Button("Downvote 👎")
        retry_button = gr.Button("Retry ♻️")
        clear_button = gr.Button("Clear History 🗑️")
    notification = gr.Markdown("")

    send_button.click(
        generate_response,
        inputs=[user_message, chat_history, messages_state],
        outputs=[user_message, chat_history, messages_state]
    )
    # for the ENTER key
    user_message.submit(
        generate_response,
        inputs=[user_message, chat_history, messages_state],
        outputs=[user_message, chat_history, messages_state]
    )

    upvote_button.click(
        upvote_click,
        inputs=[chat_history, messages_state, dataset_state],
        outputs=[dataset_state, notification]
    )

    downvote_button.click(
        downvote_click,
        inputs=[chat_history, messages_state, dataset_state],
        outputs=[dataset_state, notification]
    )

    retry_button.click(
        retry_response,
        inputs=[chat_history, messages_state],
        outputs=[chat_history, messages_state, notification]
    )

    clear_button.click(
        clear_history,
        outputs=[chat_history, messages_state]
    )

demo.launch(debug=True)


Help me with this problem: The model needs to have a more cohesive feedback that the user can provide, but it must be in the dataset format of: prompt, completion, and label. The label is only binary of true and false and the format is unable to change. But you are free to do anything to help make the model's feedback more useful for training. Feel free to add more UI elements, e.g. when downvoting or upvoting, to do this. (Have a slider allowing the user to rate usefulness, accuracy, fluency plus a comment box the user can use)