Update app.py

app.py

@@ -4,7 +4,7 @@ import streamlit as st
 import faiss
 import numpy as np
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from sentence_transformers import SentenceTransformer
 from reportlab.lib.pagesizes import A4
 from reportlab.platypus import Paragraph, SimpleDocTemplate, Spacer
@@ -23,7 +23,7 @@ age_categories = {
 }
 
 # Initialize FAISS and Sentence Transformer
-model = SentenceTransformer('all-MiniLM-L6-v2')
+embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
 
 def create_faiss_index(data):
     descriptions, age_keys = [], []
@@ -32,7 +32,7 @@ def create_faiss_index(data):
             descriptions.append(entry['description'])
             age_keys.append(int(age))  # Convert age to int
     
-    embeddings = model.encode(descriptions, convert_to_numpy=True)
+    embeddings = embedding_model.encode(descriptions, convert_to_numpy=True)
     index = faiss.IndexFlatL2(embeddings.shape[1])
     index.add(embeddings)
     return index, descriptions, age_keys
@@ -41,14 +41,14 @@ index, descriptions, age_keys = create_faiss_index(milestones)
 
 # Function to retrieve the closest milestone
 def retrieve_milestone(user_input):
-    user_embedding = model.encode([user_input], convert_to_numpy=True)
+    user_embedding = 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 model and tokenizer
-model_name = "ibm/granite-13b-chat"
+# 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 = AutoModelForCausalLM.from_pretrained(
+granite_model = AutoModelForSeq2SeqLM.from_pretrained(
     model_name, torch_dtype=torch.float16, device_map="auto"
 )
 
@@ -125,4 +125,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.")