import streamlit as st
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig


model_id = "Narrativaai/BioGPT-Large-finetuned-chatdoctor"
tokenizer = AutoTokenizer.from_pretrained("microsoft/BioGPT-Large")
model = AutoModelForCausalLM.from_pretrained(model_id)


def answer_question(prompt, temperature=0.1, top_p=0.75, top_k=40, num_beams=2, **kwargs):
    inputs = tokenizer(prompt, return_tensors="pt")
    input_ids = inputs["input_ids"].to("cpu")
    attention_mask = inputs["attention_mask"].to("cpu")
    generation_config = GenerationConfig(
        temperature=temperature, top_p=top_p, top_k=top_k, num_beams=num_beams, **kwargs
    )
    with torch.no_grad():
        generation_output = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            generation_config=generation_config,
            return_dict_in_generate=True,
            output_scores=True,
            max_new_tokens=512,
            eos_token_id=tokenizer.eos_token_id,
        )
    s = generation_output.sequences[0]
    output = tokenizer.decode(s, skip_special_tokens=True)
    return output.split(" Response:")[1]


st.set_page_config(page_title="Medical Chat Bot", page_icon=":ambulance:", layout="wide")
st.title("Medical Chat Bot")
st.caption("Talk your way to better health")

with open("./sidebar.md", "r") as sidebar_file:
    sidebar_content = sidebar_file.read()

with open("./styles.md", "r") as styles_file:
    styles_content = styles_file.read()

# Display the DDL for the selected table
st.sidebar.markdown(sidebar_content)

st.write(styles_content, unsafe_allow_html=True)

    
st.write("Please enter your question below:")

# get user input
user_input = st.text_input("You: ")

if user_input:
        # generate response
        bot_response = answer_question(f"Input: {user_input}\nResponse:")
        st.write("")
        st.write("Bot:", bot_response)