Update app.py

app.py

@@ -4,7 +4,7 @@ import streamlit as st
 import faiss
 import numpy as np
 import torch
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from sentence_transformers import SentenceTransformer
 from reportlab.lib.pagesizes import A4
 from reportlab.platypus import Paragraph, SimpleDocTemplate, Spacer
@@ -45,12 +45,20 @@ def retrieve_milestone(user_input):
     _, indices = index.search(user_embedding, 1)
     return descriptions[indices[0][0]] if indices[0][0] < len(descriptions) else "No relevant milestone found."
 
-# Load IBM Granite 3.1 model and tokenizer
-model_name = "ibm-granite/granite-3.1-8b-instruct"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-granite_model = AutoModelForSeq2SeqLM.from_pretrained(
-    model_name, torch_dtype=torch.float16, device_map="auto"
-)
+# Load IBM Granite 3.1 model and tokenizer from Hugging Face
+MODEL_NAME = "ibm-granite/granite-3.1-8b-instruct"
+
+@st.cache_resource  # Cache model to avoid reloading on every interaction
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,  
+        device_map="auto"  # Auto-select GPU/CPU
+    )
+    return tokenizer, model
+
+tokenizer, granite_model = load_model()
 
 def generate_response(user_input, child_age):
     relevant_milestone = retrieve_milestone(user_input)
@@ -61,31 +69,22 @@ def generate_response(user_input, child_age):
         "If there are any concerns, suggest steps the parents can take."
     )
 
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
+    inputs = tokenizer(prompt, return_tensors="pt").to(granite_model.device)
     output = granite_model.generate(**inputs, max_length=512)
     return tokenizer.decode(output[0], skip_special_tokens=True)
 
 # Streamlit UI Styling
 st.set_page_config(page_title="Tiny Triumphs Tracker", page_icon="👶", layout="wide")
 
-st.markdown("""
-    <style>
-        .stApp { background-color: #1e1e2e; color: #ffffff; }
-        .stTitle { text-align: center; color: #ffcc00; font-size: 36px; font-weight: bold; }
-        .stButton > button { background-color: #ffcc00; color: #000; border-radius: 5px; font-weight: bold; }
-        .stSelectbox, .stTextArea { background-color: #2e2e42; color: #ffffff; border-radius: 5px; }
-    </style>
-""", unsafe_allow_html=True)
-
-st.markdown("<h1 class='stTitle'>👶 Tiny Triumphs Tracker</h1>", unsafe_allow_html=True)
-st.markdown("Track your child's key growth milestones from birth to 5 years and detect early developmental concerns.", unsafe_allow_html=True)
+st.markdown("<h1 style='text-align:center; color:#ffcc00;'>👶 Tiny Triumphs Tracker</h1>", unsafe_allow_html=True)
+st.markdown("Track your child's key growth milestones from birth to 5 years and detect early developmental concerns.")
 
 # User selects child's age
 selected_age = st.selectbox("📅 Select child's age:", list(age_categories.keys()))
 child_age = age_categories[selected_age]
 
 # User input for traits and skills
-placeholder_text = "For example, your child might say simple words like 'mama' and 'dada' and smile when spoken to. They may grasp small objects with their fingers and show excitement during playtime."
+placeholder_text = "For example, your child might say simple words like 'mama' and 'dada' and smile when spoken to."
 user_input = st.text_area("✍️ Enter child's behavioral traits and skills:", placeholder=placeholder_text)
 
 def generate_pdf_report(ai_response):
@@ -125,4 +124,4 @@ if st.button("🔍 Analyze", help="Click to analyze the child's development mile
     with open(pdf_file, "rb") as f:
         st.download_button(label="📥 Download Progress Report", data=f, file_name="progress_report.pdf", mime="application/pdf")
 
-st.warning("⚠️ The results provided are generated by AI and should be interpreted with caution. Please consult a pediatrician for professional advice.")
+st.warning("⚠️ The results provided are generated by AI and should be interpreted with caution. Please consult a pediatrician for professional advice.")