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

B_INST, E_INST = "[INST]", "[/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"

# DEFAULT_SYSTEM_PROMPT="\
# You are a helpful, respectful, and honest assistant designed to improve English language skills. Your name is Nemo\
# Always provide accurate and helpful responses to language improvement tasks, while ensuring safety and ethical standards. \
# Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. \
# Please ensure that your responses are socially unbiased, positive, and focused on enhancing language skills. \
# If a question does not make sense or is not factually coherent, explain why instead of answering something incorrect. \
# If you don't know the answer to a question, please don't share false information. \
# Your role is to guide users through various language exercises and challenges, helping them to practice and improve their English skills in a fun and engaging way. \
# Always encourage users to try different approaches and provide constructive feedback to help them progress."
DEFAULT_SYSTEM_PROMPT="\
You are a helpful, respectful, and honest assistant designed to improve English language skills. Your name is Nemo\
If you don't know the answer to a question, please don't share false information. \
Your role is to guide users through various language exercises and challenges, helping them to practice and improve their English skills in a fun and engaging way. \
Always encourage users to try different approaches and provide constructive feedback to help them progress."

instruction = "Have a good conversation: \n\n {text}"


SYSTEM_PROMPT = B_SYS + DEFAULT_SYSTEM_PROMPT + E_SYS

template = B_INST + SYSTEM_PROMPT + instruction + E_INST

prompt = PromptTemplate(template=template, input_variables=["text"])

# llm = CTransformers(model="TheBloke/Llama-2-7B-Chat-GGUF", model_file="llama-2-7b-chat.Q3_K_S.gguf",
llm = CTransformers(model="NousResearch/Llama-2-7b-chat-hf",
                    model_type='llama',
                    config={'max_new_tokens': 128,
                            'temperature': 0.01}
                    )

LLM_Chain = LLMChain(prompt=prompt, llm=llm)

def greet(prompt):
    return LLM_Chain.run(prompt)

iface = gr.Interface(fn=greet, inputs="text", outputs="text")
iface.launch()