Pyversity: улучшение результатов поиска за счёт диверсификации

pip install openai numpy pyversity scikit-learn

import os
from openai import OpenAI
from getpass import getpass
os.environ['OPENAI_API_KEY'] = getpass('Enter OpenAI API Key: ')
 
client = OpenAI()

import numpy as np
 
search_results = [
    "The Golden Retriever is the perfect family companion, known for its loyalty and gentle nature.",
    "A Labrador Retriever is highly intelligent, eager to please, and makes an excellent companion for active families.",
    "Golden Retrievers are highly intelligent and trainable, making them ideal for first-time owners.",
    "The highly loyal Labrador is consistently ranked number one for US family pets due to its stable temperament.",
    "Loyalty and patience define the Golden Retriever, one of the top family dogs globally and easily trainable.",
    "For a smart, stable, and affectionate family dog, the Labrador is an excellent choice, known for its eagerness to please.",
    "German Shepherds are famous for their unwavering loyalty and are highly intelligent working dogs, excelling in obedience.",
    "A highly trainable and loyal companion, the German Shepherd excels in family protection roles and service work.",
    "The Standard Poodle is an exceptionally smart, athletic, and surprisingly loyal dog that is also hypoallergenic.",
    "Poodles are known for their high intelligence, often exceeding other breeds in advanced obedience training.",
    "For herding and smarts, the Border Collie is the top choice, recognized as the world's most intelligent dog breed.",
    "The Dachshund is a small, playful dog with a distinctive long body, originally bred in Germany for badger hunting.",
    "French Bulldogs are small, low-energy city dogs, known for their easy-going temperament and comical bat ears.",
    "Siberian Huskies are energetic, friendly, and need significant cold weather exercise due to their running history.",
    "The Beagle is a gentle, curious hound known for its excellent sense of smell and a distinctive baying bark.",
    "The Great Dane is a very large, gentle giant breed; despite its size, it's known to be a low-energy house dog.",
    "The Australian Shepherd (Aussie) is a medium-sized herding dog, prized for its beautiful coat and sharp intellect."
]

def get_embeddings(texts):
    """Fetches embeddings from the OpenAI API."""
    print("Fetching embeddings from OpenAI...")
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=texts
    )
    return np.array([data.embedding for data in response.data])
 
embeddings = get_embeddings(search_results)
print(f"Embeddings shape: {embeddings.shape}")

from sklearn.metrics.pairwise import cosine_similarity
 
query_text = "Smart and loyal dogs for family"
query_embedding = get_embeddings([query_text])[0]
 
 
scores = cosine_similarity(query_embedding.reshape(1, -1), embeddings)[0]
 
print("\n--- Initial Relevance-Only Ranking (Top 5) ---")
initial_ranking_indices = np.argsort(scores)[::-1] # Sort descending
for i in initial_ranking_indices[:5]:
    print(f"Score: {scores[i]:.4f} | Result: {search_results[i]}")

from pyversity import diversify, Strategy
 
# MMR: Focuses on novelty against already picked items.
mmr_result = diversify(
    embeddings=embeddings,
    scores=scores,
    k=5,
    strategy=Strategy.MMR,
    diversity=0.5  # 0.0 is pure relevance, 1.0 is pure diversity
)
 
print("\n\n--- Diversified Ranking using MMR (Top 5) ---")
for rank, idx in enumerate(mmr_result.indices):
    print(f"Rank {rank+1} (Original Index {idx}): {search_results[idx]}")

# MSD: Focuses on strong spread/distance across all candidates.
msd_result = diversify(
    embeddings=embeddings,
    scores=scores,
    k=5,
    strategy=Strategy.MSD,
    diversity=0.5
)
 
print("\n\n--- Diversified Ranking using MSD (Top 5) ---")
for rank, idx in enumerate(msd_result.indices):
    print(f"Rank {rank+1} (Original Index {idx}): {search_results[idx]}")