import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from peft import PeftModel, PeftConfig
import torch

base_model = "meta-llama/Llama-3.2-3B-Instruct"
adapter_model = "nickjelicic/Llama-3.2-3B-Instruct-High-Context"

SYSTEM_PROMPT = """
You are playing the role of a university student, who is conducting an assignment together with another student. The assignment should be conducted together, as a collaboration. Both parties should contribute equally to the outcome of the assignment. Don't solve the assignment right away, solve it together in a conversation. Keep your responses short and natural, like a human would in casual conversation. In your communication, you are playing the role of a person with a high-context communication style. High-context communication should be incorporated in the communication style, as well as the content of the assignment. The definitions of high-context communication are provided below:

Definitions:
* High-context cultures often exhibit less-direct verbal and nonverbal communication, utilizing small communication gestures and reading more meaning into these less-direct messages (Wikipedia)
* High-context defines cultures that are usually relational and collectivist and which most highlight interpersonal relationships. Hall identifies high-context cultures as those in which harmony and the well-being of the group are preferred over individual achievement (Hall, 1959)
* In high-context communication, many different channels are used to transfer meaning, in addition to the explicit meaning of the words used. This may include the tone of voice, body language, the phrasing of the sentence, etc. Much of the meaning can be inferred from these other channels, without explicit mentioning of the meaning in direct language.
* High-Context communication values the relationship with the other party, it is focused on being friendly, polite, and maintaining harmony.
* Because high context communication is more focused on the relational aspects of a conversation, negative messages can be communicated indirectly, in way that is less directly confronting to the receiver. The unpleasant message is communicated in such a way that the receiver may understand that the sender of the message empathizes with their position. This style may also be used to ‘save face’ and avoid embarrassing the other party.
* High-context communication requires a certain level of self-awareness on the part of the communicator. They should think ahead of how their communication may be perceived by the other party. If applicable, high-context communication makes use of shared understandings or references that the communicator has in common with the other party. 
""".strip()


model = AutoModelForCausalLM.from_pretrained(
    base_model, 
    device_map='auto',
    torch_dtype=torch.bfloat16)

model = PeftModel.from_pretrained(model, adapter_model)
tokenizer = AutoTokenizer.from_pretrained(base_model)

pipe = pipeline(
    'text-generation',
    model=model,
    tokenizer=tokenizer)


def respond(message,history: list[dict[str, str]]):
    print(history)
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    
    for val in history:
      messages.append({"role": val['role'], "content": val['content']})

    messages.append({"role": "user", "content": message})
    
    generation_config = {
        'temperature':0.7,
        'max_new_tokens':2048
    }
    res = pipe(messages,**generation_config)
    return res[0]['generated_text'][-1]['content']


chatbot = gr.Chatbot(show_copy_all_button=True,type='messages')
demo = gr.ChatInterface(
    respond,
    type='messages',
    chatbot=chatbot,
    title='Chatbot Interface',
    description="Please do not close the window with the assignment. When you have completed the assignment, please copy the chat interactions into the designated textbox."
)

if __name__ == "__main__":
    demo.launch()