Update app.py

app.py

@@ -1,12 +1,13 @@
 import streamlit as st
 import requests
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import json
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 # Load model and tokenizer
 @st.cache_resource
 def load_model():
-    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
-    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
+    model = AutoModelForSeq2SeqLM.from_pretrained("shreyanshjha0709/watch-description-generator")
+    tokenizer = AutoTokenizer.from_pretrained("shreyanshjha0709/watch-description-generator")
     return model, tokenizer
 
 model, tokenizer = load_model()
@@ -34,7 +35,6 @@ selected_brand = st.selectbox("Select a Brand", ["Select"] + brands)
 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)
 
@@ -56,14 +56,12 @@ if selected_brand != "Select":
                 "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)
+            inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
+            outputs = model.generate(**inputs, max_length=150, num_return_sequences=1)
             
             # Decode generated text
             description = tokenizer.decode(outputs[0], skip_special_tokens=True)
@@ -71,5 +69,38 @@ if selected_brand != "Select":
             # Display the result
             st.write("### Generated Description")
             st.write(description)
+            
+            # Display watch details
+            st.write("### Watch Details")
+            st.json(json.dumps(watch_data, indent=2))
 else:
     st.warning("Please select a brand.")
+
+# Add some information about the app
+st.sidebar.title("About")
+st.sidebar.info(
+    "This app uses a fine-tuned AI model to generate descriptions for watches. "
+    "Select a brand and a watch to get started. The model will generate a unique "
+    "description based on the watch's attributes."
+)
+
+# Add a footer
+st.markdown(
+    """
+    <style>
+    .footer {
+        position: fixed;
+        left: 0;
+        bottom: 0;
+        width: 100%;
+        background-color: #f1f1f1;
+        color: black;
+        text-align: center;
+    }
+    </style>
+    <div class="footer">
+        <p>Developed with ❤️ by Your Name</p>
+    </div>
+    """,
+    unsafe_allow_html=True
+)