Update app.py

app.py

@@ -1,54 +1,56 @@
 import streamlit as st
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-import requests
-from bs4 import BeautifulSoup
-import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
 
-# Load the pretrained model and tokenizer
-MODEL_NAME = "google/flan-t5-large"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
+# 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")
+    return model, tokenizer
 
-# Function to generate text description
-def generate_description(input_text):
-    inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
-    outputs = model.generate(inputs["input_ids"], max_length=200)
-    description = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return description
+model, tokenizer = load_model()
 
-# Function to scrape brand website
-def scrape_brand_website(brand_name, sku):
-    search_url = f"https://www.google.com/search?q={brand_name}+{sku}"
-    headers = {
-        "User-Agent": "Mozilla/5.0"
-    }
-    response = requests.get(search_url, headers=headers)
-    soup = BeautifulSoup(response.text, 'html.parser')
-    
-    # For demonstration, just fetch the first URL (modify as per your requirement)
-    link = soup.find('a')['href']
-    return link
+# Streamlit interface
+st.title("Watch Description Generator")
 
-# Streamlit App
-st.title("Watch Description Generator & Brand Scout")
+# 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"
+]
 
-# Inputs for watch details
-brand_name = st.text_input("Enter Brand Name")
+# 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 Watch SKU")
+sku = st.text_input("Enter SKU")
 
 # Button to generate description
 if st.button("Generate Description"):
-    if brand_name and watch_name and sku:
-        # Generate description using the model
-        input_text = f"Watch Name: {watch_name}, Brand: {brand_name}, SKU: {sku}"
-        generated_description = generate_description(input_text)
-        st.write("Generated Description:")
-        st.write(generated_description)
-
-        # Scrape brand website for relevant information
-        st.write("Fetching Brand Website...")
-        brand_link = scrape_brand_website(brand_name, sku)
-        st.write(f"Brand Website: {brand_link}")
+    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.warning("Please fill in all fields.")
+        st.error("Please enter all details (brand, watch name, and SKU)")