Update app.py

app.py

@@ -3,9 +3,8 @@ import json
 import streamlit as st
 import faiss
 import numpy as np
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
 from sentence_transformers import SentenceTransformer
+from transformers import pipeline
 from reportlab.lib.pagesizes import A4
 from reportlab.platypus import Paragraph, SimpleDocTemplate, Spacer
 from reportlab.lib.styles import getSampleStyleSheet
@@ -16,23 +15,32 @@ with open('milestones.json', 'r') as f:
 
 # Age categories for dropdown selection
 age_categories = {
-    "Up to 2 months": 2, "Up to 4 months": 4, "Up to 6 months": 6,
-    "Up to 9 months": 9, "Up to 1 year": 12, "Up to 15 months": 15,
-    "Up to 18 months": 18, "Up to 2 years": 24, "Up to 30 months": 30,
-    "Up to 3 years": 36, "Up to 4 years": 48, "Up to 5 years": 60
+    "Up to 2 months": 2,
+    "Up to 4 months": 4,
+    "Up to 6 months": 6,
+    "Up to 9 months": 9,
+    "Up to 1 year": 12,
+    "Up to 15 months": 15,
+    "Up to 18 months": 18,
+    "Up to 2 years": 24,
+    "Up to 30 months": 30,
+    "Up to 3 years": 36,
+    "Up to 4 years": 48,
+    "Up to 5 years": 60
 }
 
 # Initialize FAISS and Sentence Transformer
-embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+model = SentenceTransformer('all-MiniLM-L6-v2')
 
 def create_faiss_index(data):
-    descriptions, age_keys = [], []
+    descriptions = []
+    age_keys = []
     for age, categories in data.items():
         for entry in categories:
             descriptions.append(entry['description'])
-            age_keys.append(int(age))  # Convert age to int
+            age_keys.append(int(age))
     
-    embeddings = embedding_model.encode(descriptions, convert_to_numpy=True)
+    embeddings = model.encode(descriptions, convert_to_numpy=True)
     index = faiss.IndexFlatL2(embeddings.shape[1])
     index.add(embeddings)
     return index, descriptions, age_keys
@@ -41,85 +49,67 @@ index, descriptions, age_keys = create_faiss_index(milestones)
 
 # Function to retrieve the closest milestone
 def retrieve_milestone(user_input):
-    user_embedding = embedding_model.encode([user_input], convert_to_numpy=True)
+    user_embedding = model.encode([user_input], convert_to_numpy=True)
     _, 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 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()
+# Initialize IBM Granite Model
+ibm_model = pipeline("text-generation", model="ibm-granite", max_length=512)
 
 def generate_response(user_input, child_age):
     relevant_milestone = retrieve_milestone(user_input)
-    prompt = (
-        f"The child is {child_age} months old. Based on the given traits: {user_input}, "
-        f"determine whether the child is meeting expected milestones. "
-        f"Relevant milestone: {relevant_milestone}. "
-        "If there are any concerns, suggest steps the parents can take."
-    )
-
-    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)
+    prompt = (f"The child is {child_age} months old. Based on the given traits: {user_input}, "
+              f"determine whether the child is meeting expected milestones. "
+              f"Relevant milestone: {relevant_milestone}. "
+              "If there are any concerns, suggest steps the parents can take. ")
+    response = ibm_model(prompt)
+    return response[0]['generated_text']
 
 # 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; }
+    </style>
+""", 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.")
+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)
 
 # 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."
+placeholder_text = "Describe your child's behavior and skills."
 user_input = st.text_area("✍️ Enter child's behavioral traits and skills:", placeholder=placeholder_text)
 
 def generate_pdf_report(ai_response):
     pdf_file = "progress_report.pdf"
     doc = SimpleDocTemplate(pdf_file, pagesize=A4)
     styles = getSampleStyleSheet()
-    
-    elements = [
-        Paragraph("Child Development Progress Report", styles['Title']),
-        Spacer(1, 12),
-        Paragraph("Development Insights:", styles['Heading2']),
-        Spacer(1, 10)
-    ]
-
-    for part in ai_response.split('\n'):
+    elements = []
+    elements.append(Paragraph("Child Development Progress Report", styles['Title']))
+    elements.append(Spacer(1, 12))
+    elements.append(Paragraph("Development Insights:", styles['Heading2']))
+    elements.append(Spacer(1, 10))
+    response_parts = ai_response.split('\n')
+    for part in response_parts:
         part = part.strip().lstrip('0123456789.- ')
         if part:
             elements.append(Paragraph(f"• {part}", styles['Normal']))
             elements.append(Spacer(1, 5))
-
-    disclaimer = ("This report is AI-generated and is for informational purposes only. "
-                  "It should not be considered a substitute for professional medical advice. "
-                  "Always consult a qualified pediatrician for expert guidance on your child's development.")
+    disclaimer = "This report is AI-generated and is for informational purposes only. "
     elements.append(Spacer(1, 12))
     elements.append(Paragraph(disclaimer, styles['Italic']))
-
     doc.build(elements)
     return pdf_file
 
 if st.button("🔍 Analyze", help="Click to analyze the child's development milestones"):
     ai_response = generate_response(user_input, child_age)
-    
     st.subheader("📊 Development Insights:")
     st.markdown(f"<div style='background-color:#44475a; color:#ffffff; padding: 15px; border-radius: 10px;'>{ai_response}</div>", unsafe_allow_html=True)
-    
     pdf_file = generate_pdf_report(ai_response)
     with open(pdf_file, "rb") as f:
         st.download_button(label="📥 Download Progress Report", data=f, file_name="progress_report.pdf", mime="application/pdf")