Add script to generate and save embeddings for models

- Implemented `generate_embeddings.py` to load embedding models and generate embeddings for emotion and event dictionaries.
- Added functionality to save generated embeddings as pickle files in the 'embeddings' directory.
- Included error handling and logging for better debugging and tracking of the embedding generation process.

.gitignore

@@ -0,0 +1,58 @@
+# Python cache files
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+
+# Distribution / packaging
+dist/
+build/
+*.egg-info/
+
+# Virtual environments
+venv/
+env/
+ENV/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# VS Code
+.vscode/
+*.code-workspace
+
+# PyCharm
+.idea/
+
+# Logs
+*.log
+logs/
+
+# OS specific files
+.DS_Store
+Thumbs.db
+desktop.ini
+
+# Environment variables
+.env
+.env.local
+
+# Temporary files
+*.swp
+*.swo
+*~
+.temp/
+
+# NOTE: We're keeping the embeddings/*.pkl files since they're pre-generated
+# for faster startup. They're managed by Git LFS as specified in .gitattributes.
+
+# Gradio specific
+gradio_cached_examples/
+flagged/
+
+# Local development files
+.jupyter/
+.local/
+.bash_history
+.python_history

app.py

@@ -6,36 +6,155 @@ This module handles the Gradio interface and application setup.
 import gradio as gr
 from utils import logger
 from emoji_processor import EmojiProcessor
+from config import EMBEDDING_MODELS
 
 class EmojiMashupApp:
     def __init__(self):
         """Initialize the Gradio application."""
         logger.info("Initializing Emoji Mashup App")
-        self.processor = EmojiProcessor()
+        self.processor = EmojiProcessor(model_key="mpnet", use_cached_embeddings=True)  # Default to mpnet
         self.processor.load_emoji_dictionaries()
         
+    def create_model_dropdown_choices(self):
+        """Create formatted choices for the model dropdown.
+        
+        Returns:
+            List of formatted model choices
+        """
+        return [
+            f"{key} ({info['size']}) - {info['notes']}" 
+            for key, info in EMBEDDING_MODELS.items()
+        ]
+        
+    def handle_model_change(self, dropdown_value, use_cached_embeddings):
+        """Handle model selection change from dropdown.
+        
+        Args:
+            dropdown_value: Selected value from dropdown
+            use_cached_embeddings: Whether to use cached embeddings
+            
+        Returns:
+            Status message about model change
+        """
+        # Extract model key from dropdown value (first word before space)
+        model_key = dropdown_value.split()[0] if dropdown_value else "mpnet"
+        
+        # Update processor cache setting
+        self.processor.use_cached_embeddings = use_cached_embeddings
+        
+        if model_key in EMBEDDING_MODELS:
+            success = self.processor.switch_model(model_key)
+            if success:
+                cache_status = "using cached embeddings" if use_cached_embeddings else "computing fresh embeddings"
+                return f"Switched to {model_key} model ({cache_status}): {EMBEDDING_MODELS[model_key]['notes']}"
+            else:
+                return f"Failed to switch to {model_key} model"
+        else:
+            return f"Unknown model: {model_key}"
+            
+    def process_with_model(self, model_selection, text, use_cached_embeddings):
+        """Process text with selected model.
+        
+        Args:
+            model_selection: Selected model from dropdown
+            text: User input text
+            use_cached_embeddings: Whether to use cached embeddings
+            
+        Returns:
+            Tuple of (emotion emoji, event emoji, mashup image)
+        """
+        # Extract model key from dropdown value (first word before space)
+        model_key = model_selection.split()[0] if model_selection else "mpnet"
+        
+        # Update processor cache setting
+        self.processor.use_cached_embeddings = use_cached_embeddings
+        
+        if model_key in EMBEDDING_MODELS:
+            self.processor.switch_model(model_key)
+            
+        # Process text with current model
+        return self.processor.sentence_to_emojis(text)
+        
     def create_interface(self):
         """Create and configure the Gradio interface.
         
         Returns:
             Gradio Interface object
         """
-        return gr.Interface(
-            fn=self.processor.sentence_to_emojis,
-            inputs=gr.Textbox(lines=2, placeholder="Type a sentence..."),
-            outputs=[
-                gr.Text(label="Top Emotion Emoji"),
-                gr.Text(label="Top Event Emoji"),
-                gr.Image(label="Mashup Emoji")
-            ],
-            title="Sentence → Emoji Mashup",
-            description="Get the top emotion and event emoji from your sentence, and view the mashup!",
-            examples=[
-                ["I feel so happy today!"],
-                ["I'm really angry right now"],
-                ["Feeling tired after a long day"]
-            ]
-        )
+        with gr.Blocks(title="Sentence → Emoji Mashup") as interface:
+            gr.Markdown("# Sentence → Emoji Mashup")
+            gr.Markdown("Get the top emotion and event emoji from your sentence, and view the mashup!")
+            
+            with gr.Row():
+                with gr.Column(scale=3):
+                    # Model selection dropdown
+                    model_dropdown = gr.Dropdown(
+                        choices=self.create_model_dropdown_choices(),
+                        value=self.create_model_dropdown_choices()[0],  # Default to first model (mpnet)
+                        label="Embedding Model",
+                        info="Select the model used for text-emoji matching"
+                    )
+                    
+                    # Cache toggle
+                    cache_toggle = gr.Checkbox(
+                        label="Use cached embeddings",
+                        value=True,
+                        info="When enabled, embeddings will be saved to and loaded from disk"
+                    )
+                    
+                    # Text input
+                    text_input = gr.Textbox(
+                        lines=2, 
+                        placeholder="Type a sentence...",
+                        label="Your message"
+                    )
+                    
+                    # Process button
+                    submit_btn = gr.Button("Generate Emoji Mashup", variant="primary")
+                
+                with gr.Column(scale=2):
+                    # Model info display
+                    model_info = gr.Textbox(
+                        value=f"Using mpnet model (using cached embeddings): {EMBEDDING_MODELS['mpnet']['notes']}",
+                        label="Model Info", 
+                        interactive=False
+                    )
+                    
+                    # Output displays
+                    emotion_out = gr.Text(label="Top Emotion Emoji")
+                    event_out = gr.Text(label="Top Event Emoji")
+                    mashup_out = gr.Image(label="Mashup Emoji")
+            
+            # Set up event handlers
+            model_dropdown.change(
+                fn=self.handle_model_change,
+                inputs=[model_dropdown, cache_toggle],
+                outputs=[model_info]
+            )
+            
+            cache_toggle.change(
+                fn=self.handle_model_change,
+                inputs=[model_dropdown, cache_toggle],
+                outputs=[model_info]
+            )
+            
+            submit_btn.click(
+                fn=self.process_with_model,
+                inputs=[model_dropdown, text_input, cache_toggle],
+                outputs=[emotion_out, event_out, mashup_out]
+            )
+            
+            # Examples
+            gr.Examples(
+                examples=[
+                    ["I feel so happy today!"],
+                    ["I'm really angry right now"],
+                    ["Feeling tired after a long day"]
+                ],
+                inputs=text_input
+            )
+            
+        return interface
         
     def run(self, share=True):
         """Launch the Gradio application.

config.py

@@ -9,4 +9,23 @@ CONFIG = {
     "item_file": "google-emoji-kitchen-item.txt",
     "emoji_kitchen_url": "https://emojik.vercel.app/s/{emoji1}_{emoji2}",
     "default_size": 256
+}
+
+# Available embedding models
+EMBEDDING_MODELS = {
+    "mpnet": {
+        "id": "all-mpnet-base-v2",
+        "size": "110M",
+        "notes": "Balanced, great general-purpose model"
+    },
+    "gte": {
+        "id": "thenlper/gte-large",
+        "size": "335M",
+        "notes": "Context-rich, good for emotion & nuance"
+    },
+    "bge": {
+        "id": "BAAI/bge-large-en-v1.5",
+        "size": "350M",
+        "notes": "Tuned for ranking & high-precision similarity"
+    }
 }

embeddings/BAAI_bge-large-en-v1.5_emotion.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5456af7ceaa04bdc28b9b125e317eaebf503c60b6937f006b54c595850c3830a
+size 463549

embeddings/BAAI_bge-large-en-v1.5_event.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c32a321359bd0a197e906c85731003294a835e7590c655043e6e9ebdfa607de9
+size 2238733

embeddings/all-mpnet-base-v2_emotion.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6db3183f80970f30c7dee0cf846c832b5505890a071c1af8009f6ff452083f7c
+size 348852

embeddings/all-mpnet-base-v2_event.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:65d434cb2cdd1034e494a87d354345e67bfd25a90a44247cfa3406dc100334c0
+size 1684668

embeddings/thenlper_gte-large_emotion.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8b5472bf5008613f76ac06738fa55c91ff2fd6ae7472c9a1f739d210b5f2f0e
+size 463549

embeddings/thenlper_gte-large_event.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:81961955fad517b578deb1969c8e84594fe92c8eed32d6b43f85e804f5214b82
+size 2238733

emoji_processor.py

@@ -7,23 +7,36 @@ from sklearn.metrics.pairwise import cosine_similarity
 import requests
 from PIL import Image
 from io import BytesIO
+import os
 
-from config import CONFIG
-from utils import logger, kitchen_txt_to_dict
+from config import CONFIG, EMBEDDING_MODELS
+from utils import (logger, kitchen_txt_to_dict, 
+                  save_embeddings_to_pickle, load_embeddings_from_pickle,
+                  get_embeddings_pickle_path)
 
 class EmojiProcessor:
-    def __init__(self, model_name=CONFIG["model_name"]):
+    def __init__(self, model_name=None, model_key=None, use_cached_embeddings=True):
         """Initialize the emoji processor with the specified model.
         
         Args:
-            model_name: Name of the sentence transformer model to use
+            model_name: Direct name of the sentence transformer model to use
+            model_key: Key from EMBEDDING_MODELS to use (takes precedence over model_name)
+            use_cached_embeddings: Whether to use cached embeddings from pickle files
         """
+        # Get model name from the key if provided
+        if model_key and model_key in EMBEDDING_MODELS:
+            model_name = EMBEDDING_MODELS[model_key]['id']
+        elif not model_name:
+            model_name = CONFIG["model_name"]
+            
         logger.info(f"Loading model: {model_name}")
         self.model = SentenceTransformer(model_name)
+        self.current_model_name = model_name
         self.emotion_dict = {}
         self.event_dict = {}
         self.emotion_embeddings = {}
         self.event_embeddings = {}
+        self.use_cached_embeddings = use_cached_embeddings
 
     def load_emoji_dictionaries(self, emotion_file=CONFIG["emotion_file"], item_file=CONFIG["item_file"]):
         """Load emoji dictionaries from text files.
@@ -36,10 +49,95 @@ class EmojiProcessor:
         self.emotion_dict = kitchen_txt_to_dict(emotion_file)
         self.event_dict = kitchen_txt_to_dict(item_file)
         
-        # Precompute embeddings
-        logger.info("Computing embeddings for emoji dictionaries")
-        self.emotion_embeddings = {emoji: self.model.encode(desc) for emoji, desc in self.emotion_dict.items()}
-        self.event_embeddings = {emoji: self.model.encode(desc) for emoji, desc in self.event_dict.items()}
+        # Load or compute embeddings
+        self._load_or_compute_embeddings()
+    
+    def _load_or_compute_embeddings(self):
+        """Load embeddings from pickle files if available, otherwise compute them."""
+        if self.use_cached_embeddings:
+            # Try to load emotion embeddings
+            emotion_pickle_path = get_embeddings_pickle_path(self.current_model_name, "emotion")
+            loaded_emotion_embeddings = load_embeddings_from_pickle(emotion_pickle_path)
+            
+            # Try to load event embeddings
+            event_pickle_path = get_embeddings_pickle_path(self.current_model_name, "event")
+            loaded_event_embeddings = load_embeddings_from_pickle(event_pickle_path)
+            
+            # Check if we need to compute any embeddings
+            compute_emotion = loaded_emotion_embeddings is None
+            compute_event = loaded_event_embeddings is None
+            
+            if not compute_emotion:
+                # Verify all emoji keys are present in loaded embeddings
+                for emoji in self.emotion_dict.keys():
+                    if emoji not in loaded_emotion_embeddings:
+                        logger.info(f"Cached emotion embeddings missing emoji: {emoji}, will recompute")
+                        compute_emotion = True
+                        break
+                        
+                if not compute_emotion:
+                    self.emotion_embeddings = loaded_emotion_embeddings
+            
+            if not compute_event:
+                # Verify all emoji keys are present in loaded embeddings
+                for emoji in self.event_dict.keys():
+                    if emoji not in loaded_event_embeddings:
+                        logger.info(f"Cached event embeddings missing emoji: {emoji}, will recompute")
+                        compute_event = True
+                        break
+                        
+                if not compute_event:
+                    self.event_embeddings = loaded_event_embeddings
+            
+            # Compute any missing embeddings
+            if compute_emotion:
+                logger.info(f"Computing emotion embeddings for model: {self.current_model_name}")
+                self.emotion_embeddings = {emoji: self.model.encode(desc) for emoji, desc in self.emotion_dict.items()}
+                # Save for future use
+                save_embeddings_to_pickle(self.emotion_embeddings, emotion_pickle_path)
+            
+            if compute_event:
+                logger.info(f"Computing event embeddings for model: {self.current_model_name}")
+                self.event_embeddings = {emoji: self.model.encode(desc) for emoji, desc in self.event_dict.items()}
+                # Save for future use
+                save_embeddings_to_pickle(self.event_embeddings, event_pickle_path)
+        else:
+            # Compute embeddings without caching
+            logger.info("Computing embeddings for emoji dictionaries (no caching)")
+            self.emotion_embeddings = {emoji: self.model.encode(desc) for emoji, desc in self.emotion_dict.items()}
+            self.event_embeddings = {emoji: self.model.encode(desc) for emoji, desc in self.event_dict.items()}
+    
+    def switch_model(self, model_key):
+        """Switch to a different embedding model.
+        
+        Args:
+            model_key: Key from EMBEDDING_MODELS to use
+            
+        Returns:
+            True if model was switched successfully, False otherwise
+        """
+        if model_key not in EMBEDDING_MODELS:
+            logger.error(f"Unknown model key: {model_key}")
+            return False
+            
+        model_name = EMBEDDING_MODELS[model_key]['id']
+        if model_name == self.current_model_name:
+            logger.info(f"Model {model_key} is already loaded")
+            return True
+            
+        try:
+            logger.info(f"Switching to model: {model_name}")
+            self.model = SentenceTransformer(model_name)
+            self.current_model_name = model_name
+            
+            # Load or recompute embeddings with new model
+            if self.emotion_dict and self.event_dict:
+                self._load_or_compute_embeddings()
+                
+            return True
+        except Exception as e:
+            logger.error(f"Error switching model: {e}")
+            return False
 
     def find_top_emojis(self, embedding, emoji_embeddings, top_n=1):
         """Find top matching emojis based on cosine similarity.

generate_embeddings.py

@@ -0,0 +1,99 @@
+"""
+Utility script to pre-generate embedding pickle files for all models.
+
+This script will:
+1. Load each embedding model
+2. Generate embeddings for both emotion and event dictionaries
+3. Save the embeddings as pickle files in the 'embeddings' directory
+
+Run this script once locally to create all pickle files before uploading to the repository.
+"""
+
+import os
+from sentence_transformers import SentenceTransformer
+from tqdm import tqdm
+
+from config import CONFIG, EMBEDDING_MODELS
+from utils import (logger, kitchen_txt_to_dict, 
+                  save_embeddings_to_pickle, get_embeddings_pickle_path)
+
+def generate_embeddings_for_model(model_key, model_info):
+    """Generate and save embeddings for a specific model.
+    
+    Args:
+        model_key: Key of the model in EMBEDDING_MODELS
+        model_info: Model information dictionary
+        
+    Returns:
+        Tuple of (success_emotion, success_event)
+    """
+    model_id = model_info['id']
+    print(f"\nProcessing model: {model_key} ({model_id}) - {model_info['size']}")
+    
+    try:
+        # Load the model
+        print(f"Loading {model_key} model...")
+        model = SentenceTransformer(model_id)
+        
+        # Load emoji dictionaries
+        print("Loading emoji dictionaries...")
+        emotion_dict = kitchen_txt_to_dict(CONFIG["emotion_file"])
+        event_dict = kitchen_txt_to_dict(CONFIG["item_file"])
+        
+        if not emotion_dict or not event_dict:
+            print("Error: Failed to load emoji dictionaries")
+            return False, False
+            
+        # Generate emotion embeddings
+        print(f"Generating {len(emotion_dict)} emotion embeddings...")
+        emotion_embeddings = {}
+        for emoji, desc in tqdm(emotion_dict.items()):
+            emotion_embeddings[emoji] = model.encode(desc)
+            
+        # Generate event embeddings
+        print(f"Generating {len(event_dict)} event embeddings...")
+        event_embeddings = {}
+        for emoji, desc in tqdm(event_dict.items()):
+            event_embeddings[emoji] = model.encode(desc)
+            
+        # Save embeddings
+        emotion_pickle_path = get_embeddings_pickle_path(model_id, "emotion")
+        event_pickle_path = get_embeddings_pickle_path(model_id, "event")
+        
+        success_emotion = save_embeddings_to_pickle(emotion_embeddings, emotion_pickle_path)
+        success_event = save_embeddings_to_pickle(event_embeddings, event_pickle_path)
+        
+        return success_emotion, success_event
+    except Exception as e:
+        print(f"Error generating embeddings for model {model_key}: {e}")
+        return False, False
+
+def main():
+    """Main function to generate embeddings for all models."""
+    # Create embeddings directory if it doesn't exist
+    os.makedirs('embeddings', exist_ok=True)
+    
+    print(f"Generating embeddings for {len(EMBEDDING_MODELS)} models...")
+    
+    results = {}
+    
+    # Generate embeddings for each model
+    for model_key, model_info in EMBEDDING_MODELS.items():
+        success_emotion, success_event = generate_embeddings_for_model(model_key, model_info)
+        results[model_key] = {
+            'emotion': success_emotion,
+            'event': success_event
+        }
+    
+    # Print summary
+    print("\n=== Embedding Generation Summary ===")
+    for model_key, result in results.items():
+        status_emotion = "✓ Success" if result['emotion'] else "✗ Failed"
+        status_event = "✓ Success" if result['event'] else "✗ Failed"
+        print(f"{model_key:<10}: Emotion: {status_emotion}, Event: {status_event}")
+    
+    print("\nDone! Embedding pickle files are stored in the 'embeddings' directory.")
+    print("You can now upload these files to your repository.")
+
+if __name__ == "__main__":
+    main()

utils.py

@@ -3,6 +3,8 @@ Utility functions for the Emoji Mashup application.
 """
 
 import logging
+import os
+import pickle
 
 # Configure logging
 def setup_logging():
@@ -36,4 +38,60 @@ def kitchen_txt_to_dict(filepath):
         return emoji_dict
     except Exception as e:
         logger.error(f"Error loading emoji dictionary from {filepath}: {e}")
-        return {}
+        return {}
+
+def save_embeddings_to_pickle(embeddings, filepath):
+    """Save embeddings dictionary to a pickle file.
+    
+    Args:
+        embeddings: Dictionary of embeddings to save
+        filepath: Path to save the pickle file to
+        
+    Returns:
+        True if successful, False otherwise
+    """
+    try:
+        os.makedirs(os.path.dirname(filepath), exist_ok=True)
+        with open(filepath, 'wb') as f:
+            pickle.dump(embeddings, f)
+        logger.info(f"Saved embeddings to {filepath}")
+        return True
+    except Exception as e:
+        logger.error(f"Error saving embeddings to {filepath}: {e}")
+        return False
+
+def load_embeddings_from_pickle(filepath):
+    """Load embeddings dictionary from a pickle file.
+    
+    Args:
+        filepath: Path to load the pickle file from
+        
+    Returns:
+        Dictionary of embeddings if successful, None otherwise
+    """
+    if not os.path.exists(filepath):
+        logger.info(f"Pickle file {filepath} does not exist")
+        return None
+        
+    try:
+        with open(filepath, 'rb') as f:
+            embeddings = pickle.load(f)
+        logger.info(f"Loaded embeddings from {filepath}")
+        return embeddings
+    except Exception as e:
+        logger.error(f"Error loading embeddings from {filepath}: {e}")
+        return None
+
+def get_embeddings_pickle_path(model_id, emoji_type):
+    """Generate the path for an embeddings pickle file.
+    
+    Args:
+        model_id: ID of the embedding model
+        emoji_type: Type of emoji ('emotion' or 'event')
+        
+    Returns:
+        Path to the embeddings pickle file
+    """
+    # Create a safe filename from the model ID
+    safe_model_id = model_id.replace('/', '_').replace('\\', '_')
+    return os.path.join('embeddings', f"{safe_model_id}_{emoji_type}.pkl")