Demo: Semantic Movie Search

# Example: A movie description that's too long for our embedding model
movie_example = {
    "name": "Ex Machina",
    "year": 2014,
    "description": """Alex Garland's Ex Machina is a cerebral, slow-burning psychological 
    thriller that delves into the ethics and consequences of artificial intelligence. 
    The story begins with Caleb, a young programmer at a tech conglomerate, who wins 
    a contest to spend a week at the secluded estate of Nathan, the reclusive CEO..."""
    # This continues for 386 tokens - too long for optimal embedding!
}

from sentence_transformers import SentenceTransformer
from qdrant_client import QdrantClient, models

# Initialize components
encoder = SentenceTransformer("all-MiniLM-L6-v2")

# In-memory for demo: NO HNSW built -> queries are a full scan.
client = QdrantClient(":memory:")

# For ANN/HNSW:
# client = QdrantClient(url="http://localhost:6333")

# Create collection with three named vectors
client.create_collection(
    collection_name='movie_search',
    vectors_config={
        'fixed': models.VectorParams(size=384, distance=models.Distance.COSINE),
        'sentence': models.VectorParams(size=384, distance=models.Distance.COSINE),
        'semantic': models.VectorParams(size=384, distance=models.Distance.COSINE),
    },
)

from transformers import AutoTokenizer
from llama_index.core.node_parser import SentenceSplitter, SemanticSplitterNodeParser
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
MAX_TOKENS = 40

def fixed_size_chunks(text, size=MAX_TOKENS):
    """Fixed-size chunking: splits at exact token boundaries"""
    tokens = tokenizer.encode(text, add_special_tokens=False)
    return [
        tokenizer.decode(tokens[i:i+size], skip_special_tokens=True)
        for i in range(0, len(tokens), size)
    ]

def sentence_chunks(text):
    """Sentence-aware chunking: respects sentence boundaries"""
    splitter = SentenceSplitter(chunk_size=MAX_TOKENS, chunk_overlap=10)
    return splitter.split_text(text)

def semantic_chunks(text):
    """Semantic chunking: uses embedding similarity to find natural breaks.
    Note: still constrained by the embed model's context window (same as retrievers)."""
    from llama_index.core import Document
    
    semantic_splitter = SemanticSplitterNodeParser(
        buffer_size=1,
        breakpoint_percentile_threshold=95,
        embed_model=HuggingFaceEmbedding(model_name="sentence-transformers/all-MiniLM-L6-v2")
    )
    nodes = semantic_splitter.get_nodes_from_documents([Document(text=text)])
    return [node.text for node in nodes]

points = []
idx = 0

for movie in movies_data:  # Process each movie
    # Fixed-size chunks
    for chunk in fixed_size_chunks(movie["description"]):
        points.append(models.PointStruct(
            id=idx,
            vector={"fixed": encoder.encode(chunk).tolist()},
            payload={**movie, "chunk": chunk, "chunking": "fixed"}
        ))
        idx += 1

    # Sentence-aware chunks  
    for chunk in sentence_chunks(movie["description"]):
        points.append(models.PointStruct(
            id=idx,
            vector={"sentence": encoder.encode(chunk).tolist()},
            payload={**movie, "chunk": chunk, "chunking": "sentence"}
        ))
        idx += 1

    # Semantic chunks
    for chunk in semantic_chunks(movie["description"]):
        points.append(models.PointStruct(
            id=idx,
            vector={"semantic": encoder.encode(chunk).tolist()},
            payload={**movie, "chunk": chunk, "chunking": "semantic"}
        ))
        idx += 1

client.upload_points(collection_name='movie_search', points=points)
print(f"Uploaded {idx} vectors across three chunking strategies")

def search_and_compare(query, k=3):
    """Compare search results across all three chunking strategies"""
    print(f"Query: '{query}'\n")
    
    for strategy in ['fixed', 'sentence', 'semantic']:
        results = client.query_points(
            collection_name='movie_search',
            query=encoder.encode(query).tolist(),
            using=strategy,
            limit=k,
        )
        
        print(f"--- {strategy.upper()} CHUNKING ---")
        for i, point in enumerate(results.points, 1):
            payload = point.payload
            print(f"{i}. {payload['name']} ({payload['year']}) | Score: {point.score:.3f}")
            print(f"   Chunk: {payload['chunk'][:100]}...")
        print()

# Test with different queries
search_and_compare("alien invasion")
search_and_compare("questioning reality and existence")

Query: 'alien invasion'

--- FIXED CHUNKING ---
1. E.T. the Extra-Terrestrial (1982) | Score: 0.554
   Chunk: the film opens with a group of botanist aliens visiting earth, only to be interrupted...

--- SENTENCE CHUNKING ---  
1. E.T. the Extra-Terrestrial (1982) | Score: 0.568
   Chunk: The film opens with a group of botanist aliens visiting Earth, only to be interrupted...

--- SEMANTIC CHUNKING ---
1. Annihilation (2018) | Score: 0.440
   Chunk: Annihilation is not a traditional alien invasion story - it is a meditation on the fragility...

# Find movies about AI made after 2000
results = client.query_points(
    collection_name='movie_search',
    query=encoder.encode("artificial intelligence").tolist(),
    using="semantic",
    query_filter=models.Filter(
        must=[models.FieldCondition(key="year", range=models.Range(gte=2000))]
    ),
    limit=3
)

for point in results.points:
    print(f"{point.payload['name']} ({point.payload['year']}) | Score: {point.score:.3f}")

# Group by movie name to get unique recommendations
response = client.query_points_groups(
    collection_name='movie_search',
    query=encoder.encode("time travel and family relationships").tolist(),
    using="semantic",
    group_by="name",       # Group by movie title
    limit=3,               # Number of unique movies
    group_size=1,          # Best chunk per movie
)

for group in response.groups:
    print(f"{group.id} | Best match score: {group.hits[0].score:.3f}")