import streamlit as st
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

# Load the HealthScribe Clinical Note Generator model and tokenizer
@st.cache_resource
def load_model():
    model_name = "har1/HealthScribe-Clinical_Note_Generator"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    return model, tokenizer

model, tokenizer = load_model()

st.title("HealthScribe Clinical Note Generator")
st.write("Generate clinical notes based on input text.")

# Input section
input_text = st.text_area("Enter patient information or medical notes:", height=200)

if st.button("Generate Clinical Note"):
    if input_text.strip():
        # Tokenize and generate
        inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
        outputs = model.generate(inputs["input_ids"], max_length=512, num_beams=5, early_stopping=True)
        generated_note = tokenizer.decode(outputs[0], skip_special_tokens=True)

        # Display the result
        st.subheader("Generated Clinical Note")
        st.write(generated_note)
    else:
        st.warning("Please enter some text to generate a clinical note.")