Update app.py

app.py

@@ -2,17 +2,36 @@ import streamlit as st
 import torch
 from transformers import AutoModelForCausalLM
 import difflib
+import requests
+import os
+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")
+    if response.status_code == 200:
+        return response.json()
+    return None
+
+def save_to_firebase(model_id, structure):
+    response = requests.put(f"{FIREBASE_URL}/model_structures/{model_id}.json", data=json.dumps(structure))
+    return response.status_code == 200
 
-@st.cache_data
 def get_model_structure(model_id):
+    structure = fetch_from_firebase(model_id)
+    if structure:
+        return structure
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         torch_dtype=torch.bfloat16,
         device_map="cpu",
     )
     structure = {k: str(v.shape) for k, v in model.state_dict().items()}
+    save_to_firebase(model_id, structure)
     return structure
 
+
 def compare_structures(struct1, struct2):
     struct1_lines = [f"{k}: {v}" for k, v in struct1.items()]
     struct2_lines = [f"{k}: {v}" for k, v in struct2.items()]
@@ -53,8 +72,8 @@ st.markdown(
     <style>
     .reportview-container .main .block-container {
         max-width: 100%;
-        padding-left: 0%;
-        padding-right: 0%;
+        padding-left: 10%;
+        padding-right: 10%;
     }
     .stMarkdown {
         white-space: pre-wrap;