# ===========================
# Install
# ===========================
# pip install sentence-transformers
# (auto-installs transformers, torch, huggingface-hub)

from sentence_transformers import SentenceTransformer

# ===========================
# Load model (downloads from Hub, cached locally)
# ===========================
model = SentenceTransformer("all-MiniLM-L6-v2")
# all-MiniLM-L6-v2:
# - 22M parameters (TINY! DistilBERT=66M, BERT=110M)
# - 384-dimensional embeddings
# - Trained on 1 BILLION sentence pairs
# - Inference: ~14,000 sentences/second on GPU!
# - Best speed/quality ratio for English

print(f"Model loaded: {model}")
print(f"Max sequence length: {model.max_seq_length}")  # 256
print(f"Embedding dimension: {model.get_sentence_embedding_dimension()}")  # 384

# ===========================
# How Sentence Transformers works internally:
# ===========================
# 1. Tokenize sentence with BERT tokenizer
# 2. Pass through BERT/MiniLM model → get ALL token embeddings
# 3. POOL token embeddings into ONE sentence vector
#    → Mean pooling (average all token embeddings) ← most common
#    → CLS pooling (use [CLS] token only)
#    → Max pooling (take max across tokens)
# 4. NORMALIZE to unit length (for cosine similarity)
# 5. Return: 1 vector per sentence (384 or 768 dimensions)
#
# KEY DIFFERENCE from plain BERT:
# Model is FINE-TUNED on millions of sentence pairs
# using contrastive learning (similar pairs close, dissimilar far)
# → embedding space is MEANINGFUL for similarity!

from sentence_transformers import SentenceTransformer
import numpy as np

model = SentenceTransformer("all-MiniLM-L6-v2")

# ===========================
# 1. Single sentence
# ===========================
embedding = model.encode("Jakarta is the capital of Indonesia")
print(f"Type: {type(embedding)}")           # numpy.ndarray
print(f"Shape: {embedding.shape}")           # (384,)
print(f"First 5 values: {embedding[:5]}")    # [0.032, -0.018, ...]
print(f"Norm: {np.linalg.norm(embedding):.4f}")  # 1.0000 (normalized!)

# ===========================
# 2. Batch encoding (MUCH faster!)
# ===========================
sentences = [
    "I love machine learning",
    "Deep learning is fascinating",
    "The weather is beautiful today",
    "I enjoy artificial intelligence",
    "It's raining cats and dogs",
]

embeddings = model.encode(sentences, show_progress_bar=True, batch_size=32)
print(f"Batch shape: {embeddings.shape}")  # (5, 384)
# 5 sentences → 5 vectors, each 384 dimensions

# ===========================
# 3. GPU acceleration
# ===========================
model_gpu = SentenceTransformer("all-MiniLM-L6-v2", device="cuda")
embeddings = model_gpu.encode(sentences)  # ~14,000 sentences/sec on T4!

# ===========================
# 4. Return as PyTorch tensors
# ===========================
embeddings_pt = model.encode(sentences, convert_to_tensor=True)
print(f"Tensor device: {embeddings_pt.device}")  # cuda:0 (if GPU)

# ===========================
# 5. Speed benchmark
# ===========================
import time
big_corpus = [f"This is sentence number {i}" for i in range(10000)]
start = time.time()
_ = model.encode(big_corpus, batch_size=256, show_progress_bar=False)
elapsed = time.time() - start
print(f"10,000 sentences in {elapsed:.1f}s ({10000/elapsed:.0f} sent/sec)")
# GPU: 10,000 sentences in 0.7s (14,285 sent/sec)
# CPU: 10,000 sentences in 8.2s (1,220 sent/sec)

from sentence_transformers import SentenceTransformer, util

model = SentenceTransformer("all-MiniLM-L6-v2")

# ===========================
# 1. Pairwise similarity
# ===========================
sent_a = "I love machine learning"
sent_b = "Deep learning is my passion"
sent_c = "The weather is terrible today"

emb_a = model.encode(sent_a, convert_to_tensor=True)
emb_b = model.encode(sent_b, convert_to_tensor=True)
emb_c = model.encode(sent_c, convert_to_tensor=True)

sim_ab = util.cos_sim(emb_a, emb_b).item()
sim_ac = util.cos_sim(emb_a, emb_c).item()
sim_bc = util.cos_sim(emb_b, emb_c).item()

print(f"'{sent_a}' vs '{sent_b}': {sim_ab:.3f}")  # 0.782 (TINGGI — terkait!)
print(f"'{sent_a}' vs '{sent_c}': {sim_ac:.3f}")  # 0.094 (RENDAH — tidak terkait)
print(f"'{sent_b}' vs '{sent_c}': {sim_bc:.3f}")  # 0.051 (RENDAH — tidak terkait)

# ===========================
# 2. Similarity matrix (all pairs!)
# ===========================
sentences = [
    "I love cats",
    "I adore kittens",
    "Dogs are great pets",
    "The stock market crashed",
    "Financial markets are volatile",
]
embeddings = model.encode(sentences, convert_to_tensor=True)
sim_matrix = util.cos_sim(embeddings, embeddings)
print(f"Similarity matrix shape: {sim_matrix.shape}")  # (5, 5)

# Pretty print
for i in range(len(sentences)):
    for j in range(len(sentences)):
        print(f"{sim_matrix[i][j]:.2f}", end="  ")
    print(f"  ← {sentences[i][:25]}")
# 1.00  0.83  0.45  0.02  0.05  ← I love cats
# 0.83  1.00  0.48  0.01  0.04  ← I adore kittens
# 0.45  0.48  1.00  0.03  0.06  ← Dogs are great pets
# 0.02  0.01  0.03  1.00  0.79  ← The stock market crashed
# 0.05  0.04  0.06  0.79  1.00  ← Financial markets volatile

# PERFECT! Two clusters clearly visible:
# Cluster 1: animals (cats, kittens, dogs) — high similarity
# Cluster 2: finance (stock market, financial) — high similarity
# Cross-cluster: near zero — correctly unrelated!

# ===========================
# 3. Find most similar pair
# ===========================
pairs = util.paraphrase_mining(model, sentences, top_k=3)
for score, i, j in pairs:
    print(f"  {score:.3f}: '{sentences[i]}' ↔ '{sentences[j]}'")
# 0.831: 'I love cats' ↔ 'I adore kittens'
# 0.789: 'The stock market crashed' ↔ 'Financial markets are volatile'
# 0.478: 'I adore kittens' ↔ 'Dogs are great pets'

from sentence_transformers import CrossEncoder

# ===========================
# Cross-encoder: score PAIRS directly
# ===========================
cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

pairs = [
    ["How many people live in Jakarta?", "Jakarta has a population of 10 million."],
    ["How many people live in Jakarta?", "Jakarta is the capital of Indonesia."],
    ["How many people live in Jakarta?", "The weather in Jakarta is tropical."],
]

scores = cross_encoder.predict(pairs)
for pair, score in zip(pairs, scores):
    print(f"  {score:+.3f}: Q: {pair[0][:40]} | D: {pair[1][:40]}")
# +7.234: Q: How many people live in Jakarta?  | D: Jakarta has a population of 10 million  ← BEST!
# +1.123: Q: How many people live in Jakarta?  | D: Jakarta is the capital of Indonesia.
# -3.456: Q: How many people live in Jakarta?  | D: The weather in Jakarta is tropical.

# pip install faiss-cpu  (atau faiss-gpu untuk GPU)
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")

# ===========================
# 1. Create document corpus
# ===========================
documents = [
    "Python is a popular programming language",
    "Jakarta is the capital of Indonesia",
    "Machine learning uses data to make predictions",
    "The Eiffel Tower is in Paris, France",
    "TensorFlow is a deep learning framework",
    "Nasi goreng is a famous Indonesian dish",
    "Neural networks are inspired by the brain",
    "Mount Bromo is a volcano in East Java",
    "PyTorch was developed by Facebook AI",
    "Bali is a popular tourist destination",
]

# Encode all documents (ONE TIME — then saved!)
doc_embeddings = model.encode(documents, convert_to_numpy=True)
print(f"Embeddings shape: {doc_embeddings.shape}")  # (10, 384)

# ===========================
# 2. Build FAISS index
# ===========================
dimension = doc_embeddings.shape[1]  # 384

# Exact search (small datasets < 100k)
index = faiss.IndexFlatIP(dimension)  # Inner Product (= cosine sim for normalized vectors)
# IndexFlatIP = exact brute-force search
# IndexFlatL2 = L2 distance (use for non-normalized vectors)

# Normalize embeddings (required for cosine similarity with IndexFlatIP)
faiss.normalize_L2(doc_embeddings)
index.add(doc_embeddings)
print(f"Index size: {index.ntotal} vectors")  # 10

# ===========================
# 3. Search!
# ===========================
query = "What programming language is good for AI?"
query_embedding = model.encode([query], convert_to_numpy=True)
faiss.normalize_L2(query_embedding)

k = 3  # return top 3 results
scores, indices = index.search(query_embedding, k)

print(f"\\nQuery: '{query}'")
print(f"Top {k} results:")
for rank, (score, idx) in enumerate(zip(scores[0], indices[0])):
    print(f"  #{rank+1} (sim={score:.3f}): {documents[idx]}")
# #1 (sim=0.623): Python is a popular programming language
# #2 (sim=0.541): TensorFlow is a deep learning framework
# #3 (sim=0.502): Machine learning uses data to make predictions

# ===========================
# 4. For LARGE datasets (1M+ docs): use approximate index
# ===========================
# nlist = 100  # number of Voronoi cells
# quantizer = faiss.IndexFlatIP(dimension)
# index = faiss.IndexIVFFlat(quantizer, dimension, nlist)
# index.train(doc_embeddings)  # train Voronoi cells
# index.add(doc_embeddings)
# index.nprobe = 10  # search 10 nearest cells (speed/accuracy tradeoff)
# → 1M docs: ~5ms per query (vs 50ms for brute-force)

# ===========================
# 5. Save/load index
# ===========================
faiss.write_index(index, "my_search_index.faiss")
loaded_index = faiss.read_index("my_search_index.faiss")

import faiss
import numpy as np
import json
from sentence_transformers import SentenceTransformer, CrossEncoder

class SemanticSearchEngine:
    """Production semantic search: bi-encoder retrieval + cross-encoder re-ranking."""

    def __init__(self, bi_model="all-MiniLM-L6-v2",
                 cross_model="cross-encoder/ms-marco-MiniLM-L-6-v2"):
        self.bi_encoder = SentenceTransformer(bi_model)
        self.cross_encoder = CrossEncoder(cross_model)
        self.documents = []
        self.index = None
        self.dim = self.bi_encoder.get_sentence_embedding_dimension()

    def index_documents(self, documents):
        """Encode documents and build FAISS index."""
        self.documents = documents
        embeddings = self.bi_encoder.encode(documents, convert_to_numpy=True,
                                             show_progress_bar=True, batch_size=256)
        faiss.normalize_L2(embeddings)
        self.index = faiss.IndexFlatIP(self.dim)
        self.index.add(embeddings)
        print(f"Indexed {self.index.ntotal} documents")

    def search(self, query, top_k=5, rerank_top=20):
        """Two-stage search: retrieve + re-rank."""
        # Stage 1: Bi-encoder retrieval (fast!)
        q_emb = self.bi_encoder.encode([query], convert_to_numpy=True)
        faiss.normalize_L2(q_emb)
        scores, indices = self.index.search(q_emb, rerank_top)

        candidates = [(self.documents[idx], score)
                      for idx, score in zip(indices[0], scores[0]) if idx >= 0]

        # Stage 2: Cross-encoder re-ranking (accurate!)
        pairs = [[query, doc] for doc, _ in candidates]
        rerank_scores = self.cross_encoder.predict(pairs)

        # Sort by cross-encoder score
        results = sorted(zip(candidates, rerank_scores),
                         key=lambda x: x[1], reverse=True)[:top_k]

        return [{"document": doc, "bi_score": float(bi_s),
                 "rerank_score": float(re_s)}
                for (doc, bi_s), re_s in results]

    def save(self, path):
        faiss.write_index(self.index, f"{path}/index.faiss")
        with open(f"{path}/docs.json", "w") as f:
            json.dump(self.documents, f)

# ===========================
# Use it!
# ===========================
engine = SemanticSearchEngine()
engine.index_documents([
    "Python is a popular programming language for data science",
    "Jakarta is the capital and largest city of Indonesia",
    "TensorFlow and PyTorch are deep learning frameworks",
    "Machine learning models learn patterns from data",
    "Indonesia has over 17,000 islands",
    # ... add thousands of documents!
])

results = engine.search("What language is best for AI?", top_k=3)
for i, r in enumerate(results):
    print(f"  #{i+1} (rerank={r['rerank_score']:.2f}): {r['document'][:60]}...")

from sentence_transformers import SentenceTransformer, InputExample, losses
from torch.utils.data import DataLoader

# ===========================
# 1. Prepare training data (pairs + similarity score)
# ===========================
train_examples = [
    InputExample(texts=["I love cats", "I adore kittens"], label=0.9),
    InputExample(texts=["I love cats", "Stock market crashed"], label=0.0),
    InputExample(texts=["Python is great", "I enjoy coding in Python"], label=0.85),
    InputExample(texts=["Jakarta is busy", "The capital is crowded"], label=0.8),
    # ... hundreds/thousands of pairs
]

# ===========================
# 2. Fine-tune!
# ===========================
model = SentenceTransformer("all-MiniLM-L6-v2")  # start from pre-trained
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)
train_loss = losses.CosineSimilarityLoss(model)

model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=3,
    warmup_steps=100,
    output_path="./my-embedding-model",
    show_progress_bar=True,
)
# → Domain-specific embedding model! Push to Hub if you want.

# ===========================
# 3. Alternative: fine-tune with NLI data (triplets)
# ===========================
# anchor, positive, negative
triplet_examples = [
    InputExample(texts=["I love cats", "Kittens are adorable", "Stock prices fell"]),
    # ...
]
# train_loss = losses.TripletLoss(model)
# or losses.MultipleNegativesRankingLoss(model) ← BEST for retrieval!

from sentence_transformers import SentenceTransformer
from transformers import pipeline
import faiss, numpy as np

# ===========================
# 1. Setup retriever + generator
# ===========================
retriever = SentenceTransformer("all-MiniLM-L6-v2")
generator = pipeline("text2text-generation", model="google/flan-t5-small", device=0)

# ===========================
# 2. Index knowledge base
# ===========================
knowledge_base = [
    "Jakarta has a population of 10.56 million people in 2024.",
    "Indonesia declared independence on August 17, 1945.",
    "Mount Bromo is an active volcano in East Java, Indonesia.",
    "Python was created by Guido van Rossum in 1991.",
    "TensorFlow was developed by Google Brain team.",
]

kb_embeddings = retriever.encode(knowledge_base, convert_to_numpy=True)
faiss.normalize_L2(kb_embeddings)
index = faiss.IndexFlatIP(kb_embeddings.shape[1])
index.add(kb_embeddings)

# ===========================
# 3. RAG function
# ===========================
def rag_answer(question, top_k=2):
    # Retrieve
    q_emb = retriever.encode([question], convert_to_numpy=True)
    faiss.normalize_L2(q_emb)
    scores, indices = index.search(q_emb, top_k)
    context = " ".join([knowledge_base[i] for i in indices[0]])

    # Generate
    prompt = f"Answer based on this context: {context}\n\nQuestion: {question}\nAnswer:"
    result = generator(prompt, max_length=100)
    return result[0]["generated_text"], context

answer, ctx = rag_answer("What is the population of Jakarta?")
print(f"Answer: {answer}")
# "10.56 million people" ← grounded in context! ✅