import os
os.environ["HUGGINGFACEHUB_API_TOKEN"] = ""

from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain.memory import ConversationBufferMemory
from langchain_community.llms import HuggingFaceHub

template = """You are a friendly chatbot engaging in a conversation with a human.

Previous conversation:
{chat_history}

New human question: {question}
Response:"""

def get_pipeline(model_name):
    llm = HuggingFaceHub(
        repo_id=model_name,
        task="text-generation",
        model_kwargs={
            "max_new_tokens": 250,
            "top_k": 30,
            "temperature": 0.1,
            "repetition_penalty": 1.03,
        },
    )
    return llm


chatbot = get_pipeline("mistralai/Mistral-7B-Instruct-v0.2")
memory = ConversationBufferMemory(memory_key="chat_history")
prompt_template = PromptTemplate.from_template(template)
conversation = LLMChain(llm=chatbot, prompt=prompt_template, verbose=True, memory=memory)

while True:
    question = input("You: ")
    response = conversation({"question": question})
    print("-" * 50)
    print(response)
    print(response["text"])
    print("-" * 50)
    print()