import streamlit as st
from transformers import AutoTokenizer, T5ForConditionalGeneration
from peft import PeftModel, PeftConfig

@st.cache_resource
def load_model():
    base_model = T5ForConditionalGeneration.from_pretrained("google-t5/t5-small")
    tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-small")
    model = PeftModel.from_pretrained(base_model, "./")
    return tokenizer, model

tokenizer, model = load_model()

st.title("🧬 Symptom to Drug Recommendation (T5 LoRA)")

symptom_input = st.text_area("Enter Patient Symptoms:", height=150)

if st.button("Generate Treatment Plan"):
    if symptom_input.strip():
        input_text = f"symptom: {symptom_input} </s>"
        inputs = tokenizer.encode(input_text, return_tensors="pt", truncation=True)
        outputs = model.generate(inputs, max_length=100, num_beams=4, early_stopping=True)
        prediction = tokenizer.decode(outputs[0], skip_special_tokens=True)
        st.success("Predicted Medication:")
        st.write(prediction)
    else:
        st.warning("Please enter some symptoms.")