from langchain.llms import CTransformers
from langchain import PromptTemplate, LLMChain

# prepare the template we will use when prompting the AI
template = """
<s>[INST] You are a helpful, respectful and honest Medical assistant that answers
questions about the Malawian public health processes, case definitions and guidelines. 
Answer the question below exactly in few words.
{question} [/INST] </s>
"""
 
prompt = PromptTemplate(template=template, input_variables=["question"])

# load the language model
config = {'max_new_tokens': 100, 'temperature': 0}

llm = CTransformers(model='Jayem-11/OpenPipe_mistral-ft-optimized.gguf',
                    model_file="OpenPipe_mistral_optimized_merged.gguf", 
                    config=config)

def answer(query):
    llm_chain = LLMChain(prompt=prompt, llm=llm)
    return llm_chain.run({"question":query})