import streamlit as st
from langchain.llms import HuggingFaceHub
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate

# Function to return the response
def generate_answer(query):
    llm = HuggingFaceHub(
        repo_id="google/flan-t5-xxl",
        model_kwargs={"temperature": 0.7, "max_length": 64, "max_new_tokens": 512}
    )

    template = """Question: {question}
         
    Answer: Let's think step by step."""
            
    prompt = PromptTemplate(template=template, input_variables=["query"])
    llm_chain = LLMChain(prompt=prompt, llm=llm)
    result = llm_chain.run(query)
    return result

# App UI starts here
st.set_page_config(page_title="Doctor Assistant Demo", page_icon=":robot:")
st.header("Doctor Assistant Demo")

# Gets User Input
def get_text():
    input_text = st.text_input("You: ", key="input")
    return input_text

user_input = get_text()
response = generate_answer(user_input)

submit = st.button("Generate")

# If the button clicked
if submit:
    st.subheader("Doctor's Response:")
    st.write(response)