Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
+import json
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
 # Load model and tokenizer
 @st.cache_resource
@@ -10,47 +11,64 @@ def load_model():
 
 model, tokenizer = load_model()
 
-# Streamlit interface
+# Load the JSON file
+@st.cache_data
+def load_json():
+    with open('watches_data.json') as f:  # Your JSON file name
+        data = json.load(f)
+    return data
+
+data = load_json()
+
+# Extract unique brands
+brands = sorted(list(set([item["brand"] for item in data])))
+
+# Streamlit UI
 st.title("Watch Description Generator")
 
-# List of watch brands from the provided image (excluding "View All")
-brands = [
-    "Rolex", "Alpina", "Angelus", "Armin Strom", "Arnold & Son", "Baume & Mercier",
-    "Bell & Ross", "Bianchet", "Bovet", "Breitling", "Bremont", "BVLGARI",
-    "Carl F. Bucherer", "Chronoswiss", "CIGA Design", "Corum", "Cuervo y Sobrinos", 
-    "CVSTOS", "Czapek", "Doxa", "Eberhard & Co.", "Edox", "Ernest Borel", "Eterna", 
-    "Favre Leuba", "Frederique Constant", "Gerald Charles", "Girard-Perregaux", 
-    "Graham", "Grand Seiko", "H. Moser & Cie.", "Hublot", "Ikepod", "IWC Schaffhausen", 
-    "Jacob & Co.", "Jaeger-LeCoultre", "Junghans", "Laurent Ferrier", "Longines", 
-    "Louis Erard", "Louis Moinet", "Luminox", "Maserati", "Maurice Lacroix", 
-    "MeisterSinger", "Mühle-Glashütte", "Nivada Grenchen", "Nomos Glashutte", 
-    "NORQAIN", "Omega", "Oris", "Panerai", "Parmigiani", "Perrelet", "Rado", 
-    "Raymond Weil", "Ressence", "Seiko", "Speake-Marin", "TAG Heuer", "Tissot", 
-    "Titoni", "Trilobe", "Tudor", "Tutima Glashütte", "Ulysse Nardin", "Urwerk", 
-    "Zenith", "Zeppelin"
-]
-
-# Create dropdown for selecting the brand
-selected_brand = st.selectbox("Select a Brand", brands)
-
-# Input fields for Watch Name and SKU
-watch_name = st.text_input("Enter Watch Name")
-sku = st.text_input("Enter SKU")
-
-# Button to generate description
-if st.button("Generate Description"):
-    if watch_name and sku and selected_brand:
-        input_text = f"Brand: {selected_brand}, Watch Name: {watch_name}, SKU: {sku}"
-        
-        # Tokenize input and generate description
-        inputs = tokenizer(input_text, return_tensors="pt")
-        outputs = model.generate(**inputs)
-        
-        # Decode generated text
-        description = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Display the result
-        st.write("### Generated Description")
-        st.write(description)
-    else:
-        st.error("Please enter all details (brand, watch name, and SKU)")
+# Select brand
+selected_brand = st.selectbox("Select a Brand", ["Select"] + brands)
+
+# Filter watches and SKUs by the selected brand
+if selected_brand != "Select":
+    watches = [item["name"] for item in data if item["brand"] == selected_brand]
+    skus = [item["sku"] for item in data if item["brand"] == selected_brand]
+
+    selected_watch = st.selectbox("Select Watch Name (Optional)", ["Select"] + watches)
+    selected_sku = st.selectbox("Select SKU (Optional)", ["Select"] + skus)
+
+    # Get the selected watch data from the JSON
+    watch_data = None
+    if selected_watch != "Select":
+        watch_data = next((item for item in data if item["name"] == selected_watch), None)
+    elif selected_sku != "Select":
+        watch_data = next((item for item in data if item["sku"] == selected_sku), None)
+
+    if watch_data:
+        # Generate description based on attributes
+        if st.button("Generate Description"):
+            attributes = {
+                "brand": watch_data["brand"],
+                "name": watch_data.get("name", "Unknown Watch"),
+                "sku": watch_data.get("sku", "Unknown SKU"),
+                "features": watch_data.get("features", "Unknown Features"),
+                "casesize": watch_data.get("casesize", "Unknown Case Size"),
+                "movement": watch_data.get("movement", "Unknown Movement"),
+                "gender": watch_data.get("gender", "Unknown Gender"),
+                # Add more attributes as needed
+            }
+
+            input_text = f"Brand: {attributes['brand']}, Watch Name: {attributes['name']}, SKU: {attributes['sku']}, Features: {attributes['features']}, Case Size: {attributes['casesize']}, Movement: {attributes['movement']}, Gender: {attributes['gender']}"
+            
+            # Tokenize input and generate description
+            inputs = tokenizer(input_text, return_tensors="pt")
+            outputs = model.generate(**inputs)
+            
+            # Decode generated text
+            description = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            # Display the result
+            st.write("### Generated Description")
+            st.write(description)
+else:
+    st.warning("Please select a brand.")