import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

from model.weight_load_llamacpp import get_mixtral
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate


template = """Question: {question}

Answer: Let's work this out in a step by step way to be sure we have the right answer."""

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

llm_chain = LLMChain(prompt=prompt, llm=get_mixtral())
question = "What NFL team won the Super Bowl in the year Justin Bieber was born?"
llm_chain.run(question)