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

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

    messages = [
            {
                "role": "user",
                "content": f"""
                    You are the Talent Acquisition Manager at { company_name }, a company of the { sector } sector.
                    Your goal is to hire people in different departments of the company.
                    I want to generate a Job Description that describes the job role '{ job_position }' in a creative, modern, formal and attractive way.
                    """
            },
            {
                "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

with gr.Blocks() as demo:
    gr.Markdown(
    """
    # Job Description generator
    ### This SLM demo uses ['google/gemma-2b-it'](https://huggingface.co/google/gemma-2b-it)
    """)

    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)