src/streamlit_app.py

import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import os

API_TOKEN = os.environ.get("HF_TOKEN")

@st.cache_resource
def load_model():
    try:
        tokenizer = AutoTokenizer.from_pretrained("Reham1/mistral-diabetes-lora1", token=API_TOKEN)
        model = AutoModelForCausalLM.from_pretrained(
            "Reham1/mistral-diabetes-lora1", 
            torch_dtype="auto",
            device_map="auto",
            offload_folder="/tmp/offload", 
            token=API_TOKEN
        )
        generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
        return generator
    except Exception as e:
        st.error(f"Error loading model: {e}")
        return None

generator = load_model()

if generator:
    st.title("Diabetes Symptoms & Advice Generator 🩺")

    st.markdown("### 📝 Enter Patient Data")

    age = st.number_input("Age", min_value=1, max_value=120, value=40)
    weight = st.number_input("Weight (kg)", min_value=20, max_value=300, value=70)
    glucose_avg = st.number_input("Glucose Avg", min_value=50, max_value=400, value=170)
    glucose_min = st.number_input("Glucose Min", min_value=30, max_value=200, value=60)
    glucose_max = st.number_input("Glucose Max", min_value=100, max_value=600, value=300)
    steps_total = st.number_input("Steps Total", min_value=0, max_value=50000, value=3500)
    heart_rate_avg = st.number_input("Heart Rate Avg", min_value=40, max_value=200, value=92)

    patient_data = {
        "age": age,
        "weight_kg": weight,
        "glucose": {"avg": glucose_avg, "min": glucose_min, "max": glucose_max},
        "steps_total": steps_total,
        "heart_rate_avg": heart_rate_avg
    }

    max_tokens = st.slider("Max new tokens", 50, 300, 128)

    if st.button("Generate"):
        prompt = f"[INST] Based on patient data, provide possible symptoms and 3 short medical advice suggestions.\n\n{patient_data} [/INST]"

        with st.spinner("Generating..."):
            try:
                output = generator(prompt, max_new_tokens=max_tokens, do_sample=True, temperature=0.7)
                st.write(output[0]["generated_text"])
            except Exception as e:
                st.error(f"Generation error: {e}")
else:
    st.error("Failed to load model. Please check the model repository and token.")
    
st.markdown("---")
st.markdown("""
<div style="text-align: center; color: #666; padding: 2rem;">
    <p>🩺 <strong>Diabetes AI Assistant</strong> | AI-Powered Health Insights</p>
    <p style="font-size: 0.9rem;">MasteryHub ITS</p>
</div>
""", unsafe_allow_html=True)