Sync from GitHub Actions

README.md

@@ -1,3 +1,74 @@
----
-license: mit
----
+---
+title: MiniEmbed Product Matcher
+emoji: ""
+colorFrom: blue
+colorTo: indigo
+pinned: false
+license: mit
+library_name: generic
+tags:
+- embeddings
+- product-matching
+---
+
+# MiniEmbed: Product Matching Model
+
+This is a specialized version of **MiniEmbed**, fine-tuned exclusively for **high-accuracy product matching** (entity resolution). 
+
+Unlike general-purpose embedding models, this model is designed to determine if two product listings—often with different titles, specifications, or formatting—refer to the **exact same physical item**.
+
+## Use Case
+
+**Cross-Catalog Product Matching**
+*   **Scenario**: You have a catalog (Site A) and want to find matching products in a competitor's catalog (Site B).
+*   **Challenge**: Titles differ ("iPhone 14 128GB" vs "Apple iPhone 14 Midnight 128GB"), specs are formatted differently, and noise/distractors exist.
+*   **Solution**: This model maps semantically identical products to the same vector space, ignoring irrelevant noise while paying attention to critical specs (GB, Model Number, Color).
+
+## Interactive Demo
+
+This repository includes a **Streamlit** app to demonstrate the matching capability.
+
+To run locally:
+
+```bash
+pip install -r requirements.txt
+streamlit run app.py
+```
+
+## Model Architecture
+
+*   **Type**: Transformer Bi-Encoder (BERT-style)
+*   **Parameters**: ~10.8M (Mini)
+*   **Dimensions**: 256
+*   **Max Sequence Length**: 128 tokens
+*   **Format**: `SafeTensors` (Hugging Face ready)
+
+## Usage
+
+You can use the provided `src` library to run inference in your own Python scripts:
+
+```python
+from src.inference import EmbeddingInference
+
+# Load model from current directory
+model = EmbeddingInference.from_pretrained(".")
+
+# Define two product titles
+product_a = "Sony WH-1000XM5 Wireless Noise Canceling Headphones, Black"
+product_b = "Sony WH1000XM5/B Headphones"
+
+# Calculate similarity (0 to 1)
+score = model.similarity(product_a, product_b)
+
+if score > 0.82:
+    print(f"It's a match! (Score: {score:.4f})")
+else:
+    print(f"Different products. (Score: {score:.4f})")
+```
+
+## Automated Sync
+
+This repository is automatically synced to Hugging Face Spaces via GitHub Actions.
+
+
+MIT

demo.py

@@ -0,0 +1,216 @@
+
+import streamlit as st
+import sys
+import os
+import pandas as pd
+from pathlib import Path
+
+# Add root to sys.path
+sys.path.append(os.getcwd())
+
+from src.inference import EmbeddingInference
+
+st.set_page_config(page_title="Product Model Demo", layout="wide")
+
+st.title("Product Model Identity Verification Demo")
+st.markdown("""
+This demo showcases the **Product Model's** ability to verify if two product listings represent the same physical item.
+Key use case: **Matching Scale A (e.g., Your Catalog) vs Site B (e.g., Competitor/Marketplace)**.
+""")
+
+# Load Model
+@st.cache_resource
+def load_model():
+    model_path = "."
+    if not os.path.exists("pytorch_model.bin"):
+        return None
+    return EmbeddingInference.from_pretrained(model_path)
+
+model = load_model()
+
+if not model:
+    st.error("Model not found in `models/product`. Please train the model first.")
+    st.stop()
+
+st.success(f"Product Model Loaded! (Vocab: {len(model.tokenizer.word_to_id)}, Dim: {model.model.d_model})")
+
+# Settings
+st.sidebar.header("Settings")
+threshold = st.sidebar.slider("Match Threshold", 0.5, 1.0, 0.82, 0.01, help="Score above which products are considered a match.")
+
+# Data Input
+st.subheader("1. Input Data")
+
+# Session state for text areas (widget keys)
+if 'txt_a' not in st.session_state:
+    st.session_state.txt_a = """Apple iPhone 14 128GB Midnight
+Samsung Galaxy S23 Ultra 256GB Black
+Sony WH-1000XM5 Wireless Headphones
+Nintendo Switch OLED White
+Logitech MX Master 3S Performance Mouse
+Nike Air Force 1 '07 White
+Dyson V15 Detect Vacuum"""
+
+if 'txt_b' not in st.session_state:
+    st.session_state.txt_b = """iPhone 14 (128GB) - Midnight Black
+Samsung S23 Ultra 5G (256GB Storage)
+Sony Noise Cancelling Headphones WH1000XM5
+Nintendo Switch Console - OLED Model
+Logitech Mouse MX Master 3S
+Nike Men's Air Force 1 Sneakers
+Dyson V15 Detect Cordless Vacuum Cleaner
+Apple iPhone 13 128GB
+Samsung Galaxy S22 Ultra
+Sony WH-1000XM4
+Nintendo Switch Lite"""
+
+# Button to load large dataset
+if st.button("Load Large Benchmark Dataset (100+ items)"):
+    # Generate large dataset
+    base_products = [
+        ("iPhone 14 Pro 128GB Space Black", "Apple iPhone 14 Pro (128 GB) - Space Black", "iPhone 14 Pro Max 128GB"),
+        ("Samsung Galaxy S23 Ultra 512GB", "Samsung S23 Ultra 5G (512GB Storage)", "Samsung Galaxy S23 512GB"),
+        ("Sony WH-1000XM5 Headphones", "Sony Noise Cancelling Wireless Headphones WH1000XM5", "Sony WH-1000XM4 Headphones"),
+        ("MacBook Air M2 13-inch 256GB", "Apple MacBook Air Laptop: M2 chip, 13.6-inch, 256GB", "MacBook Pro M2 13-inch"),
+        ("Dyson V15 Detect Vacuum", "Dyson V15 Detect Cordless Vacuum Cleaner", "Dyson V12 Detect Slim"),
+        ("Logitech MX Master 3S", "Logitech Master Series MX 3S Mouse", "Logitech MX Master 3"),
+        ("Kindle Paperwhite 16GB", "Amazon Kindle Paperwhite (16 GB) - 6.8 display", "Kindle Paperwhite 8GB"),
+        ("Nintendo Switch OLED White", "Nintendo Switch – OLED Model w/ White Joy-Con", "Nintendo Switch Lite Blue"),
+        ("PlayStation 5 Console", "Sony PS5 Console Disc Edition", "PlayStation 5 Digital Edition"),
+        ("Xbox Series X", "Microsoft Xbox Series X 1TB Console", "Xbox Series S"),
+        ("AirPos Pro 2nd Gen", "Apple AirPods Pro (2nd Generation) with MagSafe", "Apple AirPods 3rd Gen"),
+        ("Fitbit Charge 6", "Fitbit Charge 6 Fitness Tracker with Google apps", "Fitbit Charge 5"),
+        ("Garmin Forerunner 265", "Garmin Forerunner 265 Running Smartwatch", "Garmin Forerunner 965"),
+        ("Yeti Rambler 20oz Tumbler", "YETI Rambler 20 oz Stainless Steel Vacuum Insulated", "Yeti Rambler 30oz"),
+        ("Stanley Quencher H2.0 40oz", "Stanley The Quencher H2.0 FlowState Tumbler 40oz", "Stanley IceFlow Flip Straw"),
+        ("Canon EOS R6 Mark II", "Canon Mirrorless Camera EOS R6 Mark II Body", "Canon EOS R5 Body"),
+        ("Nikon Z6 II Body", "Nikon Z 6II FX-Format Mirrorless Camera", "Nikon Z7 II Body"),
+        ("DJI Mini 3 Pro", "DJI Mini 3 Pro (DJI RC)", "DJI Mini 2 SE"),
+        ("GoPro HERO11 Black", "GoPro HERO11 Black - Waterproof Action Camera", "GoPro HERO10 Black"),
+        ("Razer DeathAdder V3 Pro", "Razer DeathAdder V3 Pro Wireless Gaming Mouse", "Razer Viper V2 Pro"),
+        ("Keychron K2 Pro Keyboard", "Keychron K2 Pro QMK/VIA Wireless Mechanical Keyboard", "Keychron K2 Version 2"),
+        ("Herman Miller Aeron Chair", "Herman Miller Aeron Ergonomic Office Chair", "Herman Miller Mirra 2"),
+        ("Instant Pot Duo Plus 6qt", "Instant Pot Duo Plus 9-in-1 Electric Pressure Cooker", "Instant Pot Duo 7-in-1"),
+        ("Ninja AF101 Air Fryer", "Ninja AF101 Air Fryer that Crisps, Roasts", "Ninja AF161 Max XL"),
+        ("Vitamix 5200 Blender", "Vitamix 5200 Blender Professional-Grade", "Vitamix E310 Explorean"),
+        ("Roomba j7+ Vacuum", "iRobot Roomba j7+ (7550) Self-Emptying Robot Vacuum", "Roomba i3+ EVO"),
+        ("Sonos Arc Soundbar", "Sonos Arc - The Premium Smart Soundbar", "Sonos Beam Gen 2"),
+        ("Bose QuietComfort 45", "Bose QuietComfort 45 Bluetooth Wireless Noise Cancelling", "Bose QuietComfort Earbuds II"),
+        ("iPad Air 5th Gen 64GB", "Apple iPad Air (5th Generation): M1 chip, 64GB", "iPad 10th Gen 64GB"),
+        ("Samsung T7 Shield 1TB", "Samsung T7 Shield 1TB Portable SSD", "Samsung T7 Touch 1TB"),
+        ("SanDisk Extreme 2TB SSD", "SanDisk 2TB Extreme Portable SSD", "SanDisk Extreme Pro 2TB"),
+        ("LG C3 OLED TV 65-inch", "LG 65-Inch Class C3 Series OLED evo 4K", "LG B3 OLED TV 65-inch"),
+        ("Samsung QN90C 55-inch", "Samsung 55-Inch Class Neo QLED 4K QN90C", "Samsung QN85C 55-inch"),
+        ("Google Pixel 7 Pro 128GB", "Google Pixel 7 Pro - 5G Android Phone 128GB", "Google Pixel 7 128GB"),
+        ("OnePlus 11 5G 16GB RAM", "OnePlus 11 5G | 16GB RAM+256GB", "OnePlus 10T 5G"),
+        ("Asus ROG Zephyrus G14", "ASUS Rogers Zephyrus G14 14” 165Hz Gaming Laptop", "Asus TUF Gaming F15"),
+        ("Dell XPS 15 9530", "Dell XPS 15 Laptop, 13th Gen Intel Core", "Dell Inspiron 16 Plus"),
+        ("Lenovo ThinkPad X1 Carbon Gen 11", "Lenovo ThinkPad X1 Carbon Gen 11 14 inch", "Lenovo ThinkPad T14s"),
+        ("HP Spectre x360 14", "HP Spectre x360 2-in-1 Laptop 13.5t", "HP Envy x360 15"),
+        ("Microsoft Surface Pro 9", "Microsoft Surface Pro 9 (2022), 13 2-in-1", "Microsoft Surface Laptop 5"),
+    ]
+
+    a_list = []
+    b_list = []
+    
+    # 1. Add core pairs
+    for a, b, distractor in base_products:
+        a_list.append(a)
+        b_list.append(b)
+        b_list.append(distractor)
+
+    # 2. Add algorithmic filler (increase to 35 iterations for >100 total)
+    for i in range(35):
+        a_list.append(f"Generic Widget Model X-{i+100} Pro")
+        b_list.append(f"Generic Widget Series X {i+100} Professional Edition")
+        b_list.append(f"Generic Widget Model X-{i+100} Lite") # Distractor
+        
+        a_list.append(f"Industrial Part #44-A{i}")
+        b_list.append(f"Genuine Industrial Part Number 44-A{i} Replacement")
+        b_list.append(f"Industrial Part #44-B{i}") # Distractor
+
+    import random
+    random.shuffle(b_list)
+    
+    # Update specific keys used by text_area to ensure UI refresh
+    st.session_state.txt_a = "\n".join(a_list)
+    st.session_state.txt_b = "\n".join(b_list)
+    
+    # Keep backing val updated too
+    st.session_state.site_a_val = st.session_state.txt_a
+    st.session_state.site_b_val = st.session_state.txt_b
+    
+    st.success(f"Loaded {len(a_list)} items with hard negatives!")
+    st.rerun()
+
+col1, col2 = st.columns(2)
+
+with col1:
+    st.markdown("**Site A (Your Catalog)**")
+    # Use key to bind to session state
+    site_a_text = st.text_area("One product per line", key="txt_a", height=300)
+
+with col2:
+    st.markdown("**Site B (Competitor/Marketplace)**")
+    site_b_text = st.text_area("One product per line", key="txt_b", height=300)
+
+# Process
+if st.button("Run Comparison", type="primary"):
+    site_a = [x.strip() for x in site_a_text.split('\n') if x.strip()]
+    site_b = [x.strip() for x in site_b_text.split('\n') if x.strip()]
+    
+    if not site_a or not site_b:
+        st.warning("Please provide data for both sites.")
+        st.stop()
+        
+    st.subheader("2. Matching Results")
+    
+    results = []
+    
+    progress_bar = st.progress(0)
+    
+    for i, product_a in enumerate(site_a):
+        # Search for best match
+        matches = model.search(product_a, site_b, top_k=1)
+        
+        if matches:
+            best = matches[0]
+            score = best['score']
+            match_product = best['text']
+            is_match = score >= threshold
+            
+            results.append({
+                "Site A Product": product_a,
+                "Best Match (Site B)": match_product,
+                "Confidence": score,
+                "Status": "Match" if is_match else "Different"
+            })
+        else:
+            results.append({
+                "Site A Product": product_a,
+                "Best Match (Site B)": "No candidate found",
+                "Confidence": 0.0,
+                "Status": "No Data"
+            })
+        
+        progress_bar.progress((i + 1) / len(site_a))
+            
+    df = pd.DataFrame(results)
+    
+    # Sort by Confidence (Desc)
+    df = df.sort_values(by="Confidence", ascending=False)
+    
+    # Styling
+    def color_status(val):
+        color = '#d4edda' if val == "Match" else '#f8d7da'
+        return f'background-color: {color}'
+
+    st.dataframe(
+        df.style.applymap(color_status, subset=['Status'])
+        .format({"Confidence": "{:.4f}"}),
+        use_container_width=True
+    )
+    
+    # Stats
+    match_count = df[df['Status'] == "Match"].shape[0]
+    st.metric("Total Matches Found", f"{match_count} / {len(site_a)}")

model.safetensors

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

requirements.txt

@@ -0,0 +1,15 @@
+# Core
+torch>=2.0.0
+numpy>=1.21.0
+tqdm>=4.64.0
+
+# Demo UI
+streamlit>=1.30.0
+plotly>=5.0.0
+
+# Optional (for clustering, CSV processing, & Benchmarking)
+scikit-learn>=1.0.0
+pandas>=2.0.0
+psutil>=5.9.0
+sentence-transformers>=2.2.0
+safetensors>=0.4.0

src/__init__.py

@@ -0,0 +1,15 @@
+"""
+MiniEmbed - Lightweight Text Embedding Model
+"""
+
+from .model import MiniTransformerEmbedding
+from .tokenizer import SimpleTokenizer
+from .inference import EmbeddingInference, EmbeddingModelManager
+
+__version__ = "1.0.0"
+__all__ = [
+    "MiniTransformerEmbedding",
+    "SimpleTokenizer", 
+    "EmbeddingInference",
+    "EmbeddingModelManager"
+]

src/inference.py

@@ -0,0 +1,338 @@
+"""
+Model Saving & Inference Module
+===================================
+Easy-to-use API for loading and running inference with the embedding model.
+"""
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+import json
+import os
+from pathlib import Path
+from typing import List, Dict, Union, Tuple
+
+from .model import MiniTransformerEmbedding
+from .tokenizer import SimpleTokenizer
+
+
+class EmbeddingModelManager:
+    """
+    Handles saving and loading the embedding model.
+    
+    Save structure:
+    model_dir/
+    ├── config.json          # Model architecture config
+    ├── model.pt             # Model weights
+    ├── tokenizer.json       # Vocabulary
+    └── training_info.json   # Training metadata (optional)
+    """
+    
+    @staticmethod
+    def save_model(
+        model: MiniTransformerEmbedding,
+        tokenizer: SimpleTokenizer,
+        save_dir: str,
+        training_info: dict = None
+    ):
+        """
+        Save model, tokenizer, and config for later use.
+        
+        Args:
+            model: Trained MiniTransformerEmbedding
+            tokenizer: SimpleTokenizer with vocabulary
+            save_dir: Directory to save model
+            training_info: Optional training metadata
+        """
+        save_dir = Path(save_dir)
+        save_dir.mkdir(parents=True, exist_ok=True)
+        
+        # 1. Save model config
+        config = {
+            'vocab_size': len(tokenizer.word_to_id),
+            'd_model': model.d_model,
+            'num_heads': model.layers[0].attention.num_heads,
+            'num_layers': len(model.layers),
+            'd_ff': model.layers[0].feed_forward.linear1.out_features,
+            'max_seq_len': model.positional_encoding.pe.size(1),
+            'pad_token_id': model.pad_token_id,
+            'size_name': save_dir.name # Use folder name as size name
+        }
+        
+        with open(save_dir / 'config.json', 'w') as f:
+            json.dump(config, f, indent=2)
+        
+        # 2. Save model weights
+        torch.save(model.state_dict(), save_dir / 'model.pt')
+        
+        # 3. Save tokenizer vocabulary
+        tokenizer.save(str(save_dir / 'tokenizer.json'))
+        
+        # 4. Save training info (optional)
+        if training_info:
+            with open(save_dir / 'training_info.json', 'w') as f:
+                json.dump(training_info, f, indent=2)
+        
+        print(f"Model saved to: {save_dir}")
+    
+    @staticmethod
+    def load_model(model_dir: str, device: str = None) -> Tuple[MiniTransformerEmbedding, SimpleTokenizer]:
+        """
+        Load model and tokenizer from directory.
+        
+        Args:
+            model_dir: Directory containing saved model
+            device: Device to load model on ('cpu', 'cuda', 'mps')
+            
+        Returns:
+            (model, tokenizer) tuple
+        """
+        model_dir = Path(model_dir)
+        
+        if device is None:
+            if torch.cuda.is_available():
+                device = 'cuda'
+            elif torch.backends.mps.is_available():
+                device = 'mps'
+            else:
+                device = 'cpu'
+        
+        # 1. Load config
+        config_path = model_dir / 'config.json'
+        
+        # If loading a checkpoint, the config might be in the 'models' folder instead
+        if not config_path.exists() and 'checkpoints' in str(model_dir):
+            potential_models_dir = Path(str(model_dir).replace('checkpoints', 'models'))
+            if (potential_models_dir / 'config.json').exists():
+                config_path = potential_models_dir / 'config.json'
+        
+        if config_path.exists():
+            with open(config_path, 'r') as f:
+                config = json.load(f)
+        else:
+            # Fallback for Product Model (Hardcoded defaults)
+            print("Warning: config.json not found. Using default Product Model configuration.")
+            config = {
+                "vocab_size": 50000,
+                "d_model": 256,
+                "num_heads": 4,
+                "num_layers": 4,
+                "d_ff": 1024,
+                "max_seq_len": 128,
+                "pad_token_id": 0
+            }
+        
+        # 2. Load tokenizer
+        tokenizer_path = model_dir / 'tokenizer.json'
+        if not tokenizer_path.exists() and 'checkpoints' in str(model_dir):
+            potential_models_dir = Path(str(model_dir).replace('checkpoints', 'models'))
+            if (potential_models_dir / 'tokenizer.json').exists():
+                tokenizer_path = potential_models_dir / 'tokenizer.json'
+
+        tokenizer = SimpleTokenizer(vocab_size=config['vocab_size'])
+        tokenizer.load(str(tokenizer_path))
+        
+        # 3. Create and load model
+        model = MiniTransformerEmbedding(
+            vocab_size=config['vocab_size'],
+            d_model=config['d_model'],
+            num_heads=config['num_heads'],
+            num_layers=config['num_layers'],
+            d_ff=config['d_ff'],
+            max_seq_len=config['max_seq_len'],
+            pad_token_id=config.get('pad_token_id', 0)
+        )
+        
+        # Try safetensors first (preferred), then pytorch_model.bin, then model.pt
+        from safetensors.torch import load_file
+        
+        safe_path = model_dir / 'model.safetensors'
+        bin_path = model_dir / 'pytorch_model.bin'
+        pt_path = model_dir / 'model.pt'
+        
+        if safe_path.exists():
+            state_dict = load_file(safe_path)
+        elif bin_path.exists():
+            state_dict = torch.load(bin_path, map_location=device, weights_only=True)
+        elif pt_path.exists():
+            state_dict = torch.load(pt_path, map_location=device, weights_only=True)
+        else:
+            raise FileNotFoundError(f"No model weights found in {model_dir}")
+            
+        # state_dict loaded, now load into model
+        model.load_state_dict(state_dict)
+        model = model.to(device)
+        model.eval()
+        
+        return model, tokenizer
+    
+    @staticmethod
+    def list_models(base_dir: str = "models") -> List[str]:
+        """
+        List available model names in the base directory.
+        
+        Returns:
+            List of directory names containing valid models
+        """
+        path = Path(base_dir)
+        if not path.exists():
+            return []
+        return sorted([d.name for d in path.iterdir() if d.is_dir() and (d / "model.pt").exists()])
+
+class EmbeddingInference:
+    """
+    High-level inference API for the embedding model.
+    
+    Usage:
+        model = EmbeddingInference.from_pretrained("./model")
+        
+        # Encode texts
+        embeddings = model.encode(["Hello world", "Machine learning"])
+        
+        # Compute similarity
+        score = model.similarity("query", "document")
+        
+        # Semantic search
+        results = model.search("python programming", documents)
+    """
+    
+    def __init__(
+        self, 
+        model: MiniTransformerEmbedding, 
+        tokenizer: SimpleTokenizer,
+        device: str = 'cpu',
+        max_length: int = 64
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = device
+        self.max_length = max_length
+        self.model.eval()
+    
+    @classmethod
+    def from_pretrained(cls, model_dir: str, device: str = None, max_length: int = 128):
+        """Load from saved model directory."""
+        model, tokenizer = EmbeddingModelManager.load_model(model_dir, device)
+        if device is None:
+            device = next(model.parameters()).device.type
+        return cls(model, tokenizer, device, max_length=max_length)
+    
+    def encode(
+        self, 
+        texts: Union[str, List[str]], 
+        batch_size: int = 32,
+        show_progress: bool = False
+    ) -> np.ndarray:
+        """
+        Encode texts to embeddings.
+        
+        Args:
+            texts: Single text or list of texts
+            batch_size: Batch size for encoding
+            show_progress: Show progress bar
+            
+        Returns:
+            numpy array of shape (n_texts, d_model)
+        """
+        if isinstance(texts, str):
+            texts = [texts]
+        
+        all_embeddings = []
+        
+        # Process in batches
+        for i in range(0, len(texts), batch_size):
+            batch_texts = texts[i:i + batch_size]
+            
+            # Tokenize
+            encodings = [
+                self.tokenizer.encode(t, self.max_length) 
+                for t in batch_texts
+            ]
+            
+            input_ids = torch.stack([e['input_ids'] for e in encodings]).to(self.device)
+            attention_mask = torch.stack([e['attention_mask'] for e in encodings]).to(self.device)
+            
+            # Encode
+            with torch.no_grad():
+                embeddings = self.model.encode(input_ids, attention_mask)
+            
+            all_embeddings.append(embeddings.cpu().numpy())
+        
+        return np.vstack(all_embeddings)
+    
+    def similarity(self, text1: str, text2: str) -> float:
+        """Compute cosine similarity between two texts."""
+        emb1 = self.encode(text1)
+        emb2 = self.encode(text2)
+        return float(np.dot(emb1[0], emb2[0]))
+    
+    def pairwise_similarity(self, texts1: List[str], texts2: List[str]) -> np.ndarray:
+        """
+        Compute pairwise similarity between two lists.
+        
+        Returns:
+            Matrix of shape (len(texts1), len(texts2))
+        """
+        emb1 = self.encode(texts1)
+        emb2 = self.encode(texts2)
+        return np.dot(emb1, emb2.T)
+    
+    def search(
+        self, 
+        query: str, 
+        documents: List[str], 
+        top_k: int = 5
+    ) -> List[Dict]:
+        """
+        Semantic search: Find most similar documents to query.
+        
+        Args:
+            query: Search query
+            documents: List of documents to search
+            top_k: Number of results to return
+            
+        Returns:
+            List of dicts with 'text', 'score', 'rank'
+        """
+        query_emb = self.encode(query)
+        doc_embs = self.encode(documents)
+        
+        # Compute similarities
+        scores = np.dot(doc_embs, query_emb.T).flatten()
+        
+        # Get top-k indices
+        top_indices = np.argsort(scores)[::-1][:top_k]
+        
+        results = []
+        for rank, idx in enumerate(top_indices, 1):
+            results.append({
+                'rank': rank,
+                'text': documents[idx],
+                'score': float(scores[idx]),
+                'index': int(idx)
+            })
+        
+        return results
+    
+    def cluster_texts(self, texts: List[str], n_clusters: int = 5) -> Dict:
+        """
+        Cluster texts by embedding similarity.
+        
+        Returns:
+            Dict with 'labels' and 'texts_by_cluster'
+        """
+        from sklearn.cluster import KMeans
+        
+        embeddings = self.encode(texts)
+        
+        kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
+        labels = kmeans.fit_predict(embeddings)
+        
+        return {
+            'labels': labels.tolist(),
+            'centroids': kmeans.cluster_centers_,
+            'texts_by_cluster': {
+                i: [texts[j] for j in range(len(texts)) if labels[j] == i]
+                for i in range(n_clusters)
+            }
+        }

src/model.py

@@ -0,0 +1,359 @@
+"""
+Mini-Transformer Embedding Model
+====================================
+A lightweight transformer encoder for generating text embeddings.
+Built from scratch using PyTorch.
+
+Architecture:
+- Token Embeddings + Sinusoidal Positional Encoding
+- N Transformer Encoder Layers (Pre-LayerNorm)
+- Multi-Head Self-Attention
+- Position-wise Feed-Forward Networks
+- Mean Pooling + L2 Normalization
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Optional
+
+
+class SinusoidalPositionalEncoding(nn.Module):
+    """
+    Sinusoidal positional encoding from "Attention Is All You Need".
+    
+    Adds position information to token embeddings using sin/cos functions
+    at different frequencies, allowing the model to understand token order.
+    """
+    
+    def __init__(self, d_model: int, max_seq_len: int = 512, dropout: float = 0.1):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        
+        # Create positional encoding matrix [max_seq_len, d_model]
+        pe = torch.zeros(max_seq_len, d_model)
+        position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)
+        
+        # Compute division term for frequencies
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
+        )
+        
+        # Apply sin to even indices, cos to odd indices
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        
+        # Add batch dimension and register as buffer (not a parameter)
+        pe = pe.unsqueeze(0)  # [1, max_seq_len, d_model]
+        self.register_buffer('pe', pe)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Tensor of shape [batch_size, seq_len, d_model]
+        Returns:
+            Tensor with positional encoding added
+        """
+        x = x + self.pe[:, :x.size(1), :]
+        return self.dropout(x)
+
+
+class MultiHeadSelfAttention(nn.Module):
+    """
+    Multi-Head Self-Attention mechanism.
+    
+    Allows the model to jointly attend to information from different
+    representation subspaces at different positions.
+    """
+    
+    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1):
+        super().__init__()
+        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
+        
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.d_k = d_model // num_heads  # Dimension per head
+        
+        # Linear projections for Q, K, V
+        self.W_q = nn.Linear(d_model, d_model)
+        self.W_k = nn.Linear(d_model, d_model)
+        self.W_v = nn.Linear(d_model, d_model)
+        
+        # Output projection
+        self.W_o = nn.Linear(d_model, d_model)
+        
+        self.dropout = nn.Dropout(dropout)
+        self.scale = math.sqrt(self.d_k)
+    
+    def forward(
+        self, 
+        x: torch.Tensor, 
+        attention_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Args:
+            x: Input tensor [batch_size, seq_len, d_model]
+            attention_mask: Optional mask [batch_size, seq_len]
+        Returns:
+            Output tensor [batch_size, seq_len, d_model]
+        """
+        batch_size, seq_len, _ = x.size()
+        
+        # Linear projections
+        Q = self.W_q(x)  # [batch, seq, d_model]
+        K = self.W_k(x)
+        V = self.W_v(x)
+        
+        # Reshape to [batch, num_heads, seq, d_k]
+        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
+        K = K.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
+        V = V.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
+        
+        # Scaled dot-product attention
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale
+        # scores: [batch, num_heads, seq, seq]
+        
+        # Apply attention mask (for padding)
+        if attention_mask is not None:
+            # Expand mask: [batch, 1, 1, seq]
+            mask = attention_mask.unsqueeze(1).unsqueeze(2)
+            scores = scores.masked_fill(mask == 0, float('-inf'))
+        
+        # Softmax and dropout
+        attn_weights = F.softmax(scores, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+        
+        # Apply attention to values
+        context = torch.matmul(attn_weights, V)
+        # context: [batch, num_heads, seq, d_k]
+        
+        # Reshape back: [batch, seq, d_model]
+        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
+        
+        # Output projection
+        output = self.W_o(context)
+        
+        return output
+
+
+class PositionwiseFeedForward(nn.Module):
+    """
+    Position-wise Feed-Forward Network.
+    
+    Two linear transformations with a GELU activation in between.
+    Applied to each position separately and identically.
+    """
+    
+    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
+        super().__init__()
+        self.linear1 = nn.Linear(d_model, d_ff)
+        self.linear2 = nn.Linear(d_ff, d_model)
+        self.dropout = nn.Dropout(dropout)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Input tensor [batch_size, seq_len, d_model]
+        Returns:
+            Output tensor [batch_size, seq_len, d_model]
+        """
+        x = self.linear1(x)
+        x = F.gelu(x)
+        x = self.dropout(x)
+        x = self.linear2(x)
+        return x
+
+
+class TransformerEncoderLayer(nn.Module):
+    """
+    Single Transformer Encoder Layer with Pre-LayerNorm.
+    
+    Components:
+    1. Multi-Head Self-Attention with residual connection
+    2. Position-wise Feed-Forward with residual connection
+    
+    Uses Pre-LayerNorm for better training stability.
+    """
+    
+    def __init__(
+        self, 
+        d_model: int, 
+        num_heads: int, 
+        d_ff: int, 
+        dropout: float = 0.1
+    ):
+        super().__init__()
+        
+        # Layer normalization
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        
+        # Sub-layers
+        self.attention = MultiHeadSelfAttention(d_model, num_heads, dropout)
+        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
+        
+        # Dropout for residual connections
+        self.dropout = nn.Dropout(dropout)
+    
+    def forward(
+        self, 
+        x: torch.Tensor, 
+        attention_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Args:
+            x: Input tensor [batch_size, seq_len, d_model]
+            attention_mask: Optional mask [batch_size, seq_len]
+        Returns:
+            Output tensor [batch_size, seq_len, d_model]
+        """
+        # Pre-norm attention block
+        normed = self.norm1(x)
+        attn_output = self.attention(normed, attention_mask)
+        x = x + self.dropout(attn_output)  # Residual connection
+        
+        # Pre-norm feed-forward block
+        normed = self.norm2(x)
+        ff_output = self.feed_forward(normed)
+        x = x + self.dropout(ff_output)  # Residual connection
+        
+        return x
+
+
+class MiniTransformerEmbedding(nn.Module):
+    """
+    Mini-Transformer Embedding Model.
+    
+    Converts variable-length text sequences into fixed-size dense vectors
+    suitable for semantic similarity, search, and clustering tasks.
+    
+    Architecture:
+    1. Token Embedding Layer (vocab → d_model)
+    2. Sinusoidal Positional Encoding
+    3. N Transformer Encoder Layers
+    4. Mean Pooling (sequence → single vector)
+    5. L2 Normalization (for cosine similarity)
+    """
+    
+    def __init__(
+        self,
+        vocab_size: int = 30000,
+        d_model: int = 256,
+        num_heads: int = 4,
+        num_layers: int = 4,
+        d_ff: int = 1024,
+        max_seq_len: int = 128,
+        dropout: float = 0.1,
+        pad_token_id: int = 0
+    ):
+        super().__init__()
+        
+        self.d_model = d_model
+        self.pad_token_id = pad_token_id
+        
+        # Token embedding
+        self.token_embedding = nn.Embedding(
+            vocab_size, d_model, padding_idx=pad_token_id
+        )
+        
+        # Positional encoding
+        self.positional_encoding = SinusoidalPositionalEncoding(
+            d_model, max_seq_len, dropout
+        )
+        
+        # Transformer encoder layers
+        self.layers = nn.ModuleList([
+            TransformerEncoderLayer(d_model, num_heads, d_ff, dropout)
+            for _ in range(num_layers)
+        ])
+        
+        # Final layer norm
+        self.final_norm = nn.LayerNorm(d_model)
+        
+        # Initialize weights
+        self._init_weights()
+    
+    def _init_weights(self):
+        """Initialize weights using Xavier/Glorot initialization."""
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, nn.Embedding):
+                nn.init.normal_(module.weight, mean=0, std=0.02)
+                if module.padding_idx is not None:
+                    nn.init.zeros_(module.weight[module.padding_idx])
+    
+    def forward(
+        self, 
+        input_ids: torch.Tensor, 
+        attention_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Forward pass through the encoder.
+        
+        Args:
+            input_ids: Token IDs [batch_size, seq_len]
+            attention_mask: Mask for padding [batch_size, seq_len]
+            
+        Returns:
+            Token-level representations [batch_size, seq_len, d_model]
+        """
+        # Token embeddings with scaling
+        x = self.token_embedding(input_ids) * math.sqrt(self.d_model)
+        
+        # Add positional encoding
+        x = self.positional_encoding(x)
+        
+        # Pass through transformer layers
+        for layer in self.layers:
+            x = layer(x, attention_mask)
+        
+        # Final layer norm
+        x = self.final_norm(x)
+        
+        return x
+    
+    def encode(
+        self, 
+        input_ids: torch.Tensor, 
+        attention_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Encode input tokens to a single embedding vector per sequence.
+        
+        Uses mean pooling over non-padded tokens, followed by L2 normalization.
+        
+        Args:
+            input_ids: Token IDs [batch_size, seq_len]
+            attention_mask: Mask for padding [batch_size, seq_len]
+            
+        Returns:
+            Normalized embeddings [batch_size, d_model]
+        """
+        # Get token-level representations
+        token_embeddings = self.forward(input_ids, attention_mask)
+        
+        # Mean pooling
+        if attention_mask is not None:
+            # Expand mask for broadcasting: [batch, seq, 1]
+            mask_expanded = attention_mask.unsqueeze(-1).float()
+            
+            # Sum of embeddings (masked)
+            sum_embeddings = torch.sum(token_embeddings * mask_expanded, dim=1)
+            
+            # Count of non-padded tokens
+            sum_mask = torch.clamp(mask_expanded.sum(dim=1), min=1e-9)
+            
+            # Mean
+            embeddings = sum_embeddings / sum_mask
+        else:
+            # Simple mean over all tokens
+            embeddings = torch.mean(token_embeddings, dim=1)
+        
+        # L2 normalization for cosine similarity
+        embeddings = F.normalize(embeddings, p=2, dim=1)
+        
+        return embeddings

src/tokenizer.py

@@ -0,0 +1,162 @@
+"""
+Simple Word-Level Tokenizer
+==============================
+A basic tokenizer for demonstration purposes.
+Converts text to token IDs with special tokens.
+"""
+
+import re
+import json
+import torch
+from typing import Dict, List, Optional
+from collections import Counter
+from tqdm import tqdm
+
+
+class SimpleTokenizer:
+    """
+    A simple word-level tokenizer with special tokens.
+    
+    Special Tokens:
+    - [PAD]: Padding token (id=0)
+    - [UNK]: Unknown token (id=1)
+    - [CLS]: Classification token (id=2)
+    - [SEP]: Separator token (id=3)
+    """
+    
+    def __init__(self, vocab_size: int = 30000):
+        self.vocab_size = vocab_size
+        
+        # Special tokens
+        self.special_tokens = {
+            '[PAD]': 0,
+            '[UNK]': 1,
+            '[CLS]': 2,
+            '[SEP]': 3,
+        }
+        
+        # Word to ID mapping
+        self.word_to_id: Dict[str, int] = dict(self.special_tokens)
+        self.id_to_word: Dict[int, str] = {v: k for k, v in self.special_tokens.items()}
+        
+        # Special token IDs
+        self.pad_token_id = 0
+        self.unk_token_id = 1
+        self.cls_token_id = 2
+        self.sep_token_id = 3
+    
+    def _tokenize(self, text: str) -> List[str]:
+        """
+        Split text into tokens (simple word-level tokenization).
+        
+        Args:
+            text: Input text string
+            
+        Returns:
+            List of tokens
+        """
+        # Lowercase and basic cleaning
+        text = text.lower().strip()
+        
+        # Simple word tokenization with punctuation handling
+        tokens = re.findall(r'\b\w+\b|[^\w\s]', text)
+        
+        return tokens
+    
+    def build_vocab(self, texts: List[str], min_freq: int = 2):
+        """
+        Build vocabulary from a list of texts.
+        
+        Args:
+            texts: List of text strings
+            min_freq: Minimum frequency for a word to be included
+        """
+        # Count word frequencies
+        word_counts = Counter()
+        
+        for text in tqdm(texts, desc="Building vocabulary"):
+            tokens = self._tokenize(text)
+            word_counts.update(tokens)
+        
+        # Sort by frequency and take top vocab_size - special_tokens
+        max_words = self.vocab_size - len(self.special_tokens)
+        
+        sorted_words = sorted(
+            word_counts.items(),
+            key=lambda x: x[1],
+            reverse=True
+        )
+        
+        # Add words to vocabulary
+        for word, count in sorted_words[:max_words]:
+            if count >= min_freq and word not in self.word_to_id:
+                idx = len(self.word_to_id)
+                self.word_to_id[word] = idx
+                self.id_to_word[idx] = word
+        
+        print(f"Vocabulary size: {len(self.word_to_id)}")
+    
+    def encode(self, text: str, max_length: int = 128) -> Dict:
+        # Tokenize
+        tokens = self._tokenize(text)
+        
+        # Convert to IDs (with CLS and SEP)
+        token_ids = [self.cls_token_id]
+        
+        for token in tokens[:max_length - 2]:  # Reserve space for CLS and SEP
+            token_id = self.word_to_id.get(token, self.unk_token_id)
+            token_ids.append(token_id)
+        
+        token_ids.append(self.sep_token_id)
+        
+        # Create attention mask
+        attention_mask = [1] * len(token_ids)
+        
+        # Pad to max_length
+        padding_length = max_length - len(token_ids)
+        token_ids.extend([self.pad_token_id] * padding_length)
+        attention_mask.extend([0] * padding_length)
+        
+        return {
+            'input_ids': torch.tensor(token_ids, dtype=torch.long),
+            'attention_mask': torch.tensor(attention_mask, dtype=torch.long)
+        }
+    
+    def decode(self, token_ids: List[int]) -> str:
+        """
+        Decode token IDs back to text.
+        
+        Args:
+            token_ids: List of token IDs
+            
+        Returns:
+            Decoded text string
+        """
+        tokens = []
+        for idx in token_ids:
+            if idx in [self.pad_token_id, self.cls_token_id, self.sep_token_id]:
+                continue
+            token = self.id_to_word.get(idx, '[UNK]')
+            tokens.append(token)
+        return ' '.join(tokens)
+    
+    def save(self, path: str):
+        """Save tokenizer vocabulary to JSON file."""
+        data = {
+            'vocab_size': self.vocab_size,
+            'word_to_id': self.word_to_id,
+        }
+        with open(path, 'w') as f:
+            json.dump(data, f, indent=2)
+    
+    def load(self, path: str):
+        """Load tokenizer vocabulary from JSON file."""
+        with open(path, 'r') as f:
+            data = json.load(f)
+        
+        self.vocab_size = data['vocab_size']
+        self.word_to_id = data['word_to_id']
+        self.id_to_word = {int(v): k for k, v in self.word_to_id.items()}
+    
+    def __len__(self) -> int:
+        return len(self.word_to_id)