# This is a sample Python script.

# Press ⌃R to execute it or replace it with your code.
# Press Double ⇧ to search everywhere for classes, files, tool windows, actions, and settings.


import gradio as gr
from transformers import pipeline
from transformers import Conversation


def chatwith_blenderbot400m():
    chatbot = pipeline(task="conversational", model="facebook/blenderbot-400M-distill")
    user_message = "What are some fun activities I can do in the winter?"
    conversation = Conversation(user_message)
    print(conversation)
    print(type(conversation))
    conversation = chatbot(conversation)
    print(conversation)
    conversation.add_message(
        {"role": "user", "content": "I would like to do outdoor activities. Which activities can I do?"})
    conversation = chatbot(conversation)
    print(conversation)


def chatwith_qwen2_1point5b_instruct():
    chatbot = pipeline(task="text-generation", model="Qwen/Qwen2-1.5B-Instruct")
    messages = [{"role": "user", "content": "What are some fun activities I can do in the winter?"}]

    messages = chatbot(messages, max_new_tokens=50)[0]["generated_text"]
    print(messages)

    messages.append({"role": "user", "content": "I would like to do outdoor activities. Which activities can I do?"})
    print(messages)
    messages = chatbot(messages, max_new_tokens=50)[0]["generated_text"]
    print(messages)


def chatwith_qwen2_1point5b_instruct(prompt, max_newtokens):
    print("Aaaaa")
    chatbot = pipeline(task="text-generation", model="Qwen/Qwen2-1.5B-Instruct")

    messages = [{"role": "user", "content": prompt}]
    messages = chatbot(messages, max_new_tokens=max_newtokens)[0]["generated_text"]
    return messages


#chatwith_blenderbot400m()
#chatwith_qwen2_1point5b_instruct()

# prompt = "What are some fun activities I can do in the winter?"
# max_newtokens = 2
# print(chatwith_qwen2_1point5b_instruct(prompt, max_newtokens))

# def greet(name, intensity):
#     return "Hello........., " + name + "!" * int(intensity)


demo = gr.Interface(
    fn=chatwith_qwen2_1point5b_instruct,
    inputs=["text", "slider"],
    outputs=["text"],
)

demo.launch(share=False)