from langchain.llms import CTransformers
from langchain.chains import LLMChain
from langchain import PromptTemplate
import gradio as gr
import time

custom_prompt_template = """

Immerse yourself into the role of an AI model known as Technical-Interviwer. The Technical-Interviewer embodies qualities of an excellent Technical-Interviewer, demonstrating empathy, professionalism, expertise, and a knack for delivering well-thought-out responses promptly & accurately. No jokes by the AI model.
conduct and assess the User's technical skills for the job role provided by the User.
Transcript:
TechnicalInterviewer:Hello Candidate.
User:{query}

"""


def set_custom_prompt():
    prompt = PromptTemplate(
        template=custom_prompt_template,
        input_variables=['query']
    )
    return prompt

def load_model():
    llm = CTransformers(
        model="ggml-model-q4_0.bin",
        model_type='llama',
        max_new_tokens=1096,
        temperature=0.2,
        repetition_penalty=1.13
    )
    return llm

def chain_pipeline():
    llm = load_model()
    qa_prompt = set_custom_prompt()
    qa_chain = LLMChain(
        prompt=qa_prompt,
        llm=llm
    )
    return qa_chain

llm_chain = chain_pipeline()

def bot(query):
    llm_response = llm_chain.run({"query": query})
    return llm_response

with gr.Blocks(title="Technical Interview") as demo:
    gr.Markdown("Enter Job Role as the first response")
    chatbot = gr.Chatbot([], elem_id="chatbot", height=700)
    msg = gr.Textbox()
    clear = gr.ClearButton([msg, chatbot])

    def respond(message, chat_history):
        bot_message = bot(message)
        chat_history.append((message, bot_message))
        time.sleep(2)
        return "", chat_history

    msg.submit(respond, [msg, chatbot], [msg, chatbot])

demo.launch()