import streamlit as st
from langchain.chains import LLMChain
from langchain_core.prompts import PromptTemplate
from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline,BitsAndBytesConfig
# quants = BitsAndBytesConfig(load_in_4bit=True)

template = ''' You are an expert Blog generator , Given the Topic , the intended audience and the maximum number of words ,
    Write a blog on the given topic

    Topic : {topic}
    Intended Audince : {role}
    Number of Words : {words}

    Strictly return the output in a markdown format.
    Return only the blog and do not provide any other information.'''


prompt = PromptTemplate(template = template,input_variables = ['topic','role','words'])


def main():
    st.title("	:fire: Professional Blog Generator 	:fire:")
    st.markdown(
        """
        <style>
        body {
            background-color: #000000;;
            color: white;
        }
        </style>
        """,
        unsafe_allow_html=True
    )

    st.sidebar.header("Input Parameters")
    role = st.sidebar.text_input("Who is this intednded for ?", "Ex - Data Scientist")
    topic = st.sidebar.text_input("On what Topic should the blog be on ?", "Ex - Machine Learning")
    word_count = st.sidebar.slider("Number of Words", min_value=50, max_value=1000, value=200, step=50)

    if st.sidebar.button("Generate Blog"):
        model_id = "mistralai/Mistral-7B-Instruct-v0.2"
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        model = AutoModelForCausalLM.from_pretrained(model_id)
        pipe = pipeline("text-generation", model=model, tokenizer=tokenizer,max_new_tokens=1000)
        hf = HuggingFacePipeline(pipeline=pipe)
        chain = LLMChain(llm=hf,prompt=prompt,verbose=True)
        aa = chain.invoke({"topic": topic,"words":word_count,"role":role})
        st.write(aa)

if __name__ == "__main__":
    main()