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

from transformers.utils import logging
logging.set_verbosity_error()

# Info about ['google/gemma-2b-it'](https://huggingface.co/google/gemma-2b-it)
chatbot = pipeline(task="conversational",
                   model="./models/google/gemma-2b-it")

# company_name = "AXA Partners"
# sector = "Insurance"
# job_position = "Tech Conversational AI Lead"

def get_job_description(company_name, sector, job_position):
    start_time = datetime.now()

    messages = [
            {
                "role": "user",
                "content": f"You are the bot director of Human Resources at { company_name }, a company of the { sector } sector. " +
                    "Your goal is to hire people in different departments of the company. " +
                    "Generate a text that describes the job in a creative, modern, formal and attractive way."
            },
            {
                "role": "assistant",
                "content": "Tell me what job role it is and I'll generate the description for you!"
            },
            {
                "role": "user",
                "content": job_position
            },
            {
                "role": "assistant",
                "content": "Description:"
            }
        ]
    
    conversation = Conversation(messages)
    conversation = chatbot(conversation, do_sample=True, max_new_tokens=600, temperature=1)
    response = conversation.messages[-1]["content"]
    
    delta_time = datetime.now() - start_time

    return response, delta_time

# response = chatbot_query("AXA Partners", "Insurance", "Tech Conversational AI Lead")
# print("\n", response)

with gr.Blocks() as demo:
    gr.Markdown(
    """
    # Job Descriptions generator
    """)

    interface = gr.Interface(
        fn = get_job_description,
        inputs = [gr.Textbox(placeholder="Enter the company name", label="Company name"),
                  gr.Textbox(placeholder="Enter the company sector", label="Company sector"),
                  gr.Textbox(placeholder="Enter the job position", label="Job position")],
        outputs = [gr.Markdown(),
               gr.Textbox(label="Elapsed time")]
    )

demo.launch(server_port=7860)
# demo.launch(share=True, server_port=7860)