adding both models for classification

app.py

@@ -19,44 +19,40 @@ sidebar_markdown = """
 Note, this demo can classify 200 items. If you didn't find what you're looking for, reach out to us on our [Community](https://join.slack.com/t/marqo-community/shared_invite/zt-2iab0260n-QJrZLUSOJYUifVxf964Gdw) and request an item to be added.
 
 ## Documentation
-📚 [Blog Post]()
+📚 [Blog Post](https://www.marqo.ai/blog/introducing-marqos-ecommerce-embedding-models)
 
-📝 [Use Case Blog Post]()
+📝 [Classification Use Case Blog Post](https://www.marqo.ai/blog/ecommerce-image-classification-with-huggingface-transformers)
+
+🔎 [Image Search Use Case Blog Post](https://www.marqo.ai/blog/how-to-build-an-ecommerce-image-search-application)
 
 ## Code
-💻 [GitHub Repo]()
+💻 [GitHub Repo](https://github.com/marqo-ai/marqo-ecommerce-embeddings)
 
-🤝 [Google Colab]()
+🤝 [Google Colab](https://colab.research.google.com/drive/1ctqDrXs_P-RIOPc9xcUF83WLdYQ0wf-8?usp=sharing)
 
-🤗 [Hugging Face Collection]()
+🤗 [Hugging Face Collection](https://huggingface.co/collections/Marqo/marqo-ecommerce-embeddings-66f611b9bb9d035a8d164fbb)
 
 ## Citation
 If you use Marqo-Ecommerce-L or Marqo-Ecommerce-B, please cite us:
 ```
-
+@software{zhu2024marqoecommembed_2024,
+    author = {Tianyu Zhu and and Jesse Clark},
+    month = oct,
+    title = {{Marqo Ecommerce Embeddings - Foundation Model for Product Embeddings}},
+    url = {https://github.com/marqo-ai/marqo-ecommerce-embeddings/},
+    version = {1.0.0},
+    year = {2024}
+    }
 ```
 """
 
-from huggingface_hub import login
-
-# Get your Hugging Face API key (ensure it is set in your environment variables)
-api_key = os.getenv("HF_API_TOKEN")
-
-if api_key is None:
-    raise ValueError("Hugging Face API key not found. Please set the 'HF_API_TOKEN' environment variable.")
-
-# Login using the token
-login(token=api_key)
-
-# Initialize the model and tokenizer
-@spaces.GPU
-def load_model(progress=gr.Progress()):
-    progress(0, "Initializing model...")
-    model_name = 'hf-hub:Marqo/marqo-ecommerce-embeddings-B'
-    model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms(model_name)
+# Function to initialize a model, preprocess, and text features
+def initialize_model(model_name, progress=gr.Progress()):
+    progress(0, f"Initializing model: {model_name}...")
+    model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms(f"hf-hub:Marqo/{model_name}")
     
     progress(0.5, "Loading tokenizer...")
-    tokenizer = open_clip.get_tokenizer(model_name)
+    tokenizer = open_clip.get_tokenizer(f"hf-hub:Marqo/{model_name}")
     
     text = tokenizer(ecommerce_items)
     
@@ -65,16 +61,23 @@ def load_model(progress=gr.Progress()):
         text_features = model.encode_text(text)
         text_features /= text_features.norm(dim=-1, keepdim=True)
     
-    progress(1.0, "Model loaded successfully!")
+    progress(1.0, f"Model {model_name} loaded successfully!")
     
     return model, preprocess_val, text_features
 
-# Load model and prepare interface
-model, preprocess_val, text_features = load_model()
+# Load L model first, followed by B model
+progress_bar = gr.Progress()
+model_l, preprocess_val_l, text_features_l = initialize_model("marqo-ecommerce-embeddings-L", progress=progress_bar)
+model_b, preprocess_val_b, text_features_b = initialize_model("marqo-ecommerce-embeddings-B", progress=progress_bar)
 
 # Prediction function
 @spaces.GPU
-def predict(image, url):
+def predict(image, url, model_name):
+    if model_name == "marqo-ecommerce-embeddings-B":
+        model, preprocess_val, text_features = model_b, preprocess_val_b, text_features_b
+    else:
+        model, preprocess_val, text_features = model_l, preprocess_val_l, text_features_l
+    
     if url:
         response = requests.get(url)
         image = Image.open(BytesIO(response.content))
@@ -133,14 +136,19 @@ with gr.Blocks(css="""
         with gr.Column(scale=2):
             input_image = gr.Image(type="pil", label="Upload Ecommerce Item Image", height=312)
             input_url = gr.Textbox(label="Or provide an image URL")
+            model_selector = gr.Dropdown(
+                choices=["marqo-ecommerce-embeddings-L", "marqo-ecommerce-embeddings-B"],
+                value="marqo-ecommerce-embeddings-L",
+                label="Select Model"
+            )
             with gr.Row():
                 predict_button = gr.Button("Classify")
                 clear_button = gr.Button("Clear")
             gr.Markdown("Or click on one of the images below to classify it:")
             gr.Examples(examples=examples, inputs=input_image)
             output_label = gr.Label(num_top_classes=6)
-            predict_button.click(predict, inputs=[input_image, input_url], outputs=[input_image, output_label])
-            clear_button.click(clear_fields, outputs=[input_image, input_url])
+            predict_button.click(predict, inputs=[input_image, input_url, model_selector], outputs=[input_image, output_label])
+            clear_button.click(clear_fields, outputs=[input_image, input_url, model_selector])
 
 # Launch the interface
-demo.launch()
+demo.launch()