editors-note / app.py
shreyanshjha0709's picture
Update app.py
caabb16 verified
raw
history blame
6.2 kB
import streamlit as st
import requests
from bs4 import BeautifulSoup
import json
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
# Load model and tokenizer
@st.cache_resource
def load_model():
model = AutoModelForSeq2SeqLM.from_pretrained("shreyanshjha0709/watch-description-generator")
tokenizer = AutoTokenizer.from_pretrained("shreyanshjha0709/watch-description-generator")
return model, tokenizer
model, tokenizer = load_model()
# Load the JSON file from a URL
@st.cache_data
def load_json_from_url(url):
response = requests.get(url)
return response.json()
# Provide your JSON URL here
json_url = "https://www.ethoswatches.com/feeds/holbox_ai.json"
data = load_json_from_url(json_url)
# Extract unique brands
brands = sorted(list(set([item["brand"] for item in data])))
# Function to scrape Ethos product description using the specified selector
def scrape_ethos_description(product_link):
try:
response = requests.get(product_link)
soup = BeautifulSoup(response.text, 'html.parser')
# Target the specific selector to extract the description
description = soup.select_one('#brand_collection > div > div.lHeight_100.spec_editorNotes > div > p:nth-child(5)')
if description:
return description.get_text(strip=True)
else:
return "No detailed description available from Ethos."
except Exception as e:
return f"Error fetching details from Ethos: {str(e)}"
# Streamlit UI
st.title("Watch Description Generator")
# 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:
# Display the image from the JSON
image_url = watch_data.get("image", None)
if image_url:
st.image(image_url, caption=f"{watch_data['name']} Image")
# Get the Ethos product link for web scraping
product_link = watch_data.get("url", None)
if product_link:
st.write(f"Fetching details from: [Product Page]({product_link})")
# Scrape Ethos product page for description
ethos_description = scrape_ethos_description(product_link)
st.write("### Ethos Product Description (Extracted)")
st.write(ethos_description)
else:
st.warning("No Ethos link available for this SKU.")
# Generate a watch description based on attributes and scraped content
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"),
"water_resistance": watch_data.get("water_resistance", "Unknown Water Resistance"),
"power_reserve": watch_data.get("power_reserve", "Unknown Power Reserve"),
"dial_color": watch_data.get("dial_color", "Unknown Dial Color"),
"strap_material": watch_data.get("strap_material", "Unknown Strap Material")
}
# Combine Ethos description and attributes into a prompt
input_text = f"""Generate a detailed 100-word description for the following watch:
Brand: {attributes['brand']}
Name: {attributes['name']}
SKU: {attributes['sku']}
Features: {attributes['features']}
Case Size: {attributes['casesize']}
Movement: {attributes['movement']}
Gender: {attributes['gender']}
Water Resistance: {attributes['water_resistance']}
Power Reserve: {attributes['power_reserve']}
Dial Color: {attributes['dial_color']}
Strap Material: {attributes['strap_material']}
Additional details from Ethos:
{ethos_description}
Description: Provide a luxurious, detailed description focusing on the craftsmanship, innovation, and design, in a style similar to high-end editorials."""
# Tokenize input and generate description
inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
outputs = model.generate(
**inputs,
max_length=150, # Output length 100-150 words
num_return_sequences=1,
temperature=0.7,
top_k=50,
top_p=0.95,
do_sample=True,
repetition_penalty=1.2 # Prevent repetition
)
# Decode generated text
description = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Display the final generated description
st.write("### Final Generated Description")
st.write(description)
else:
st.warning("Please select a brand.")
# Sidebar information
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 and additional details from the Ethos website."
)
# 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 Shreyansh Jha</p>
</div>
""",
unsafe_allow_html=True
)