Update app.py

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 import torch
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import difflib
 import requests
 import os
@@ -8,23 +8,20 @@ import json
 
 FIREBASE_URL = os.getenv("FIREBASE_URL")
 
-
-def fetch_from_firebase(model_id):
-    response = requests.get(f"{FIREBASE_URL}/model_structures/{model_id}.json")
+def fetch_from_firebase(model_id, data_type):
+    response = requests.get(f"{FIREBASE_URL}/{data_type}/{model_id}.json")
     if response.status_code == 200:
         return response.json()
     return None
 
-
-def save_to_firebase(model_id, structure):
+def save_to_firebase(model_id, data, data_type):
     response = requests.put(
-        f"{FIREBASE_URL}/model_structures/{model_id}.json", data=json.dumps(structure)
+        f"{FIREBASE_URL}/{data_type}/{model_id}.json", data=json.dumps(data)
     )
     return response.status_code == 200
 
-
 def get_model_structure(model_id) -> list[str]:
-    struct_lines = fetch_from_firebase(model_id)
+    struct_lines = fetch_from_firebase(model_id, "model_structures")
     if struct_lines:
         return struct_lines
     model = AutoModelForCausalLM.from_pretrained(
@@ -34,17 +31,22 @@ def get_model_structure(model_id) -> list[str]:
     )
     structure = {k: str(v.shape) for k, v in model.state_dict().items()}
     struct_lines = [f"{k}: {v}" for k, v in structure.items()]
-    save_to_firebase(model_id, struct_lines)
+    save_to_firebase(model_id, struct_lines, "model_structures")
     return struct_lines
 
+def get_tokenizer_vocab_size(model_id) -> int:
+    vocab_size = fetch_from_firebase(model_id, "tokenizer_vocab_sizes")
+    if vocab_size:
+        return vocab_size
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    vocab_size = tokenizer.vocab_size
+    save_to_firebase(model_id, vocab_size, "tokenizer_vocab_sizes")
+    return vocab_size
 
 def compare_structures(struct1_lines: list[str], struct2_lines: list[str]):
-    # struct1_lines = [f"{k}: {v}" for k, v in struct1.items()]
-    # struct2_lines = [f"{k}: {v}" for k, v in struct2.items()]
     diff = difflib.ndiff(struct1_lines, struct2_lines)
     return diff
 
-
 def display_diff(diff):
     left_lines = []
     right_lines = []
@@ -74,7 +76,6 @@ def display_diff(diff):
 
     return left_html, right_html, diff_found
 
-
 # Set Streamlit page configuration to wide mode
 st.set_page_config(layout="wide")
 
@@ -99,10 +100,7 @@ st.title("Model Structure Comparison Tool")
 model_id1 = st.text_input("Enter the first HuggingFace Model ID")
 model_id2 = st.text_input("Enter the second HuggingFace Model ID")
 
-if "compare_button_clicked" not in st.session_state:
-    st.session_state.compare_button_clicked = False
-
-if st.session_state.compare_button_clicked:
+if st.button("Compare Models"):
     with st.spinner('Comparing models and loading tokenizers...'):
         if model_id1 and model_id2:
             struct1 = get_model_structure(model_id1)
@@ -127,15 +125,11 @@ if st.session_state.compare_button_clicked:
                 
             # Tokenizer verification
             try:
-                tokenizer1 = AutoTokenizer.from_pretrained(model_id1)
-                tokenizer2 = AutoTokenizer.from_pretrained(model_id2)
-                st.write(f"**{model_id1} Tokenizer Vocab Size**: {tokenizer1.vocab_size}")
-                st.write(f"**{model_id2} Tokenizer Vocab Size**: {tokenizer2.vocab_size}")
+                vocab_size1 = get_tokenizer_vocab_size(model_id1)
+                vocab_size2 = get_tokenizer_vocab_size(model_id2)
+                st.write(f"**{model_id1} Tokenizer Vocab Size**: {vocab_size1}")
+                st.write(f"**{model_id2} Tokenizer Vocab Size**: {vocab_size2}")
             except Exception as e:
                 st.error(f"Error loading tokenizers: {e}")
         else:
             st.error("Please enter both model IDs.")
-        st.session_state.compare_button_clicked = False
-else:
-    if st.button("Compare Models"):
-        st.session_state.compare_button_clicked = True