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-Collectivist"

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 collectivist cultural background. Collectivism should be incorporated in the communication style, as well as the content of the assignment. The definitions of collectivism are provided below:

Definitions:
* Collectivism is a philosophy that emphasizes the connection between the individual and the community. Its overriding philosophy is based on the belief that a person's social identity and personality are largely molded by community relationships. People in collectivistic cultures are closely connected and promote harmony and unity within society as a whole. (Source: Wikipedia)
* Collectivism focuses on increasing overall group functioning as a way to create more overall happiness. (Source: Trompenaars)
* In collectivist cultures, the needs, values, and goals of the in-group take precedence over the needs, values, and goals of the individual. (Triandis, 1995).
* People in collectivist cultures tend to apply different value standards to ingroups and outgroups. (Gudykunst, 1997).
* Collectivist communication is characterized by more frequent use of ‘we’ form.
* Decisions are often referred back for verification. People ideally achieve in groups which assume joint responsibility. Vacations in organized groups or with extended family.
* When doing business in a collectivist organization, negotiators can only agree tentatively and may withdraw an undertaking after consulting with superiors. The aim is to build lasting relationships.
* Managers in collectivist organizations seek to integrate personality with authority within the group. Give attention to esprit de corps, morale and cohesiveness.
* Have low job turnover and mobility. Extol the whole group and avoid showing favoritism. 
""".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()