Spaces:
Sleeping
Sleeping
import streamlit as st | |
import requests | |
import json | |
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM | |
# Load model and tokenizer | |
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 | |
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]))) | |
# 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: | |
# Generation parameters | |
max_length = st.slider("Max Length", min_value=100, max_value=500, value=300) # Increase max length to allow for more text | |
temperature = st.slider("Temperature", min_value=0.1, max_value=1.0, value=0.7, step=0.1) | |
repetition_penalty = st.slider("Repetition Penalty", min_value=1.0, max_value=2.0, value=1.2, step=0.1) # Penalize repetitive text | |
# 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"), | |
} | |
# Format input similar to training data (adjust this based on your training data format) | |
input_text = f"""Generate a detailed 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']} | |
Description:""" | |
# Tokenize input and generate description | |
inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True) | |
outputs = model.generate( | |
**inputs, | |
max_length=max_length, | |
num_return_sequences=1, | |
temperature=temperature, | |
top_k=50, | |
top_p=0.95, | |
do_sample=True, | |
repetition_penalty=repetition_penalty # Prevent repetitive descriptions | |
) | |
# Decode generated text | |
description = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
# 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)) | |
# Display input text for debugging | |
st.write("### Input Text (for debugging)") | |
st.text(input_text) | |
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 Shreyansh Jha</p> | |
</div> | |
""", | |
unsafe_allow_html=True | |
) | |