📐 Embedding Geometry & Semantic Search Playground¶

import os
import sys
import sys
from pathlib import Path

try:
    import google.colab
    IN_COLAB = True
except ImportError:
    IN_COLAB = False

if IN_COLAB:
    print("🤖 Running in Google Colab. Setting up environment...")
    !pip install -q numpy scikit-learn matplotlib loguru
    
    # Clone repository to access custom local source modules
    if not os.path.exists("production-ai-systems"):
        !git clone https://github.com/RogerioLS/production-ai-systems.git
    
    sys.path.append("production-ai-systems")
    print("🎉 Colab environment successfully initialized!")
# Ensure project root is in the path when running locally
if not IN_COLAB:
    # Localiza dinamicamente a raiz do repositório procurando por 'pyproject.toml'
    current_dir = Path.cwd()
    root_dir = current_dir
    for parent in [current_dir] + list(current_dir.parents):
        if (parent / "pyproject.toml").exists() or (parent / "projects").exists():
            root_dir = parent
            break

    if str(root_dir) not in sys.path:
        sys.path.insert(0, str(root_dir))

    print(f"📁 Project root added to sys.path: {root_dir}")

import matplotlib.pyplot as plt
import numpy as np

from projects.a_llm_basics.src.lab_02_embeddings.embeddings import WordCategoryEmbedder
from projects.a_llm_basics.src.lab_02_embeddings.reducer import DimensionalityReducer
from projects.a_llm_basics.src.lab_02_embeddings.semantic_search import SemanticSearchEngine, cosine_similarity

print("✅ Imports loaded successfully!")

vector_a = np.array([1.0, 0.0, 0.0])
vector_b = np.array([0.0, 1.0, 0.0])  # Orthogonal
vector_c = np.array([-1.0, 0.0, 0.0]) # Opposite

print(f"Similarity A vs B (Orthogonal): {cosine_similarity(vector_a, vector_b):.4f}")
print(f"Similarity A vs C (Opposite):   {cosine_similarity(vector_a, vector_c):.4f}")
print(f"Similarity A vs A (Collinear):  {cosine_similarity(vector_a, vector_a):.4f}")

embedder = WordCategoryEmbedder(dimension=100, noise_std=0.06, seed=42)

# Look at vocabulary categories
for category, words in embedder.categories.items():
    print(f"Category: {category:<8} | Words: {words}")

word_pairs = [
    ("apple", "banana"),
    ("python", "java"),
    ("apple", "python"),
    ("soccer", "athlete")
]

for w1, w2 in word_pairs:
    v1 = embedder.embed_text(w1)
    v2 = embedder.embed_text(w2)
    sim = cosine_similarity(v1, v2)
    print(f"Similarity between '{w1}' and '{w2}': {sim:.4f}")

engine = SemanticSearchEngine(embedder=embedder)

# Index standard semantic documents
engine.add_documents([
    "apple",
    "python",
    "soccer",
    "dog"
])

# Enter your custom query words below
query = "computer"
results = engine.search(query, top_k=3)

print(f"🔎 Search Results for query: '{query}'")
for rank, (doc, score) in enumerate(results, 1):
    print(f"#{rank} Word: '{doc}' with Cosine Similarity: {score:.4f}")

# Gather all words in categories
all_words = []
all_labels = []
all_colors = []
color_map = {
    "food": "#ff7f0e",
    "tech": "#1f77b4",
    "sports": "#2ca02c",
    "animals": "#d62728"
}

for category, cat_words in embedder.categories.items():
    for w in cat_words:
        all_words.append(w)
        all_labels.append(category)
        all_colors.append(color_map[category])

vectors = embedder.embed_batch(all_words)

# Project to 2D using PCA & t-SNE
pca = DimensionalityReducer(method="pca", n_components=2)
tsne = DimensionalityReducer(method="tsne", n_components=2, perplexity=5)

coords_pca = pca.fit_transform(vectors)
coords_tsne = tsne.fit_transform(vectors)

# Plotting 2D projections
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 7))

# PCA Scatter
for cat, col in color_map.items():
    indices = [i for i, label in enumerate(all_labels) if label == cat]
    ax1.scatter(coords_pca[indices, 0], coords_pca[indices, 1], c=col, label=cat.capitalize(), edgecolors="k", s=120, alpha=0.85)
ax1.set_title("PCA Projection (2D)", fontsize=13, fontweight="bold")
ax1.grid(True, linestyle="--", alpha=0.5)
for i, w in enumerate(all_words):
    ax1.annotate(w, (coords_pca[i, 0], coords_pca[i, 1]), textcoords="offset points", xytext=(0, 6), ha="center", fontsize=8)

# t-SNE Scatter
for cat, col in color_map.items():
    indices = [i for i, label in enumerate(all_labels) if label == cat]
    ax2.scatter(coords_tsne[indices, 0], coords_tsne[indices, 1], c=col, label=cat.capitalize(), edgecolors="k", s=120, alpha=0.85)
ax2.set_title("t-SNE Projection (2D)", fontsize=13, fontweight="bold")
ax2.grid(True, linestyle="--", alpha=0.5)
for i, w in enumerate(all_words):
    ax2.annotate(w, (coords_tsne[i, 0], coords_tsne[i, 1]), textcoords="offset points", xytext=(0, 6), ha="center", fontsize=8)

plt.suptitle("Dimensionality Reduction Comparison", fontsize=16, fontweight="bold")
plt.legend(loc="upper center", bbox_to_anchor=(-0.1, -0.08), ncol=4)
plt.show()

%matplotlib inline

pca_3d = DimensionalityReducer(method="pca", n_components=3)
coords_3d = pca_3d.fit_transform(vectors)

fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(projection="3d")
ax.set_facecolor("#fcfcfc")

# Draw 3D scatter
for cat, col in color_map.items():
    indices = [i for i, label in enumerate(all_labels) if label == cat]
    ax.scatter(
        coords_3d[indices, 0], coords_3d[indices, 1], coords_3d[indices, 2], 
        c=col, label=cat.capitalize(), edgecolors="k", s=100, alpha=0.85
    )

# Annotate text
for i, w in enumerate(all_words):
    ax.text(coords_3d[i, 0], coords_3d[i, 1], coords_3d[i, 2], w, fontsize=8, fontweight="semibold")

ax.set_xlabel("PC 1")
ax.set_ylabel("PC 2")
ax.set_zlabel("PC 3")
ax.set_title("3D Embedding Projection (PCA)", fontsize=14, fontweight="bold")
plt.legend()
plt.show()