import streamlit as st
from langchain import HuggingFaceHub
import os
from langchain import PromptTemplate, LLMChain


def get_llm():

    huggingfacehub_api_token = os.environ['HF_TOKEN']

    #repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
    repo_id = "Salvatale/Llama-2-7b-chat-finetune"
    llm = HuggingFaceHub(huggingfacehub_api_token=huggingfacehub_api_token,
                        repo_id=repo_id,
                        model_kwargs={"temperature":0.5,"max_new_tokens":5000})
    return llm

def get_prompt():

    template = """
    You are a helpful AI assistant and provide the answer for the question asked politely.

    {question}

    Answer:
    """

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

    return prompt


def main():
    #uploaded_file = st.file_uploader('Choose your .pdf file', type="pdf")
    # if uploaded_file is not None:
    #     df = extract_data(uploaded_file)

    llm = get_llm()
    prompt = get_prompt()

    llm_chain = LLMChain(prompt=prompt, llm=llm, verbose=True)

    st.header("Chatbot")

    user_input = st.text_input("You: ", "")
    if user_input:
        response = str(llm_chain.run(user_input))
        raws = response.splitlines()

        new_text = '\n'.join(raws[7:])
        
        st.text_area("Chatbot:", value=new_text)

if __name__ == "__main__":
    main()