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
import os
from langchain_community.llms import HuggingFaceEndpoint
aa=''

HF_TOKEN = os.environ["HF_TOKEN"]
# 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 ?", "Data Scientist")
    topic = st.sidebar.text_input("On what Topic should the blog be on ?", "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"):
        repo_id = "google/gemma-1.1-7b-it"

        llm = HuggingFaceEndpoint(
            repo_id=repo_id, max_length=128, temperature=0.5, huggingfacehub_api_token=HF_TOKEN 
        )
        llm_chain = LLMChain(prompt=prompt, llm=llm)
        # print(llm_chain.run(question))
        aa = llm_chain.run({"topic": topic,"words":word_count,"role":role})
        st.write(aa)
        
    
    # st.write(aa)
    # st.write("Will Come here")

if __name__ == "__main__":
    main()