import streamlit as st
from langchain import PromptTemplate
from langchain.chains import LLMChain

# Import your model and template
llm_hf = HuggingFaceHub(
    repo_id="tiiuae/falcon-7b-instruct",
    model_kwargs={"temperature": 0.3}
)

restaurant_template = """
I want you to become an indian lawyer and enlist me the acts and rights with a one liner description, for every issue i'll raise, you've to only give me relevant acts and rights covering from that issue...
"""

# Define the Streamlit app
def main():
    st.title("Indian Legal Consultation")
    
    # Input for the user's issue
    description = st.text_area("Please describe your issue:")

    if st.button("Get Legal Advice"):
        if description:
            # Create a PromptTemplate
            prompt = PromptTemplate(
                input_variables=["issue"],
                template=restaurant_template,
            )
            
            # Create an LLMChain
            chain = LLMChain(llm=llm_hf, prompt=prompt)
            
            # Generate a response
            response = chain.run(description)
            
            # Display the response
            st.subheader("Legal Advice:")
            st.write(response)
        else:
            st.warning("Please provide a description of your issue.")

if __name__ == "__main__":
    main()