使用LLM嵌入构建简单的语义搜索引擎

import pandas as pd
import json
from pydantic import BaseModel, Field
from openai import OpenAI
from google.colab import userdata
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from sklearn.preprocessing import StandardScaler

from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from sklearn.neighbors import NearestNeighbors

print("Loading dataset...")
dataset = load_dataset("ag_news", split="train[:1000]") # Extract the text column into a Python list
documents = dataset["text"]
print(f"Loaded {len(documents)} documents.")
print(f"Sample: {documents[0][:100]}...")

print("Loading embedding model...")
model = SentenceTransformer("all-MiniLM-L6-v2") # Convert text documents into numerical vector embeddings
print("Encoding documents (this may take a few seconds)...")
document_embeddings = model.encode(documents, show_progress_bar=True)
print(f"Created {document_embeddings.shape[0]} embeddings.")

search_engine = NearestNeighbors(n_neighbors=5, metric="cosine")
search_engine.fit(document_embeddings)
print("Search engine is ready!")

def semantic_search(query, top_k=3):
    # Embed the incoming search query
    query_embedding = model.encode([query])

    # Retrieve the closest matches
    distances, indices = search_engine.kneighbors(query_embedding, n_neighbors=top_k)

    print(f"\n🔍 Query: '{query}'")
    print("-" * 50)
    for i in range(top_k):
        doc_idx = indices[0][i]
        # Convert cosine distance to similarity (1 - distance)
        similarity = 1 - distances[0][i]
        print(f"Result {i+1} (Similarity: {similarity:.4f})")
        print(f"Text: {documents[int(doc_idx)][:150]}...\n")

semantic_search("Wall street and stock market trends")
semantic_search("Space exploration and rocket launches")

🔍 Query: 'Wall street and stock market trends'
--------------------------------------------------
Result 1 (Similarity: 0.6258)
Text: Stocks Higher Despite Soaring Oil Prices NEW YORK - Wall Street shifted higher Monday as bargain hunters shrugged off skyrocketing oil prices and boug...
Result 2 (Similarity: 0.5586)
Text: Stocks Sharply Higher on Dip in Oil Prices NEW YORK - A drop in oil prices and upbeat outlooks from Wal-Mart and Lowe's prompted new bargain-hunting o...
Result 3 (Similarity: 0.5459)
Text: Strategies for a Sideways Market (Reuters) Reuters - The bulls and the bears are in this together, scratching their heads and wondering what's t...

🔍 Query: 'Space exploration and rocket launches'
--------------------------------------------------
Result 1 (Similarity: 0.5803)
Text: Redesigning Rockets: NASA Space Propulsion Finds a New Home (SPACE.com) SPACE.com - While the exploration of the Moon and other planets in our solar s...
Result 2 (Similarity: 0.5008)
Text: Canadian Team Joins Rocket Launch Contest (AP) AP - The #36;10 million competition to send a private manned rocket into space started looking more li...
Result 3 (Similarity: 0.4724)
Text: The Next Great Space Race: SpaceShipOne and Wild Fire to Go For the Gold (SPACE.com) SPACE.com - A piloted rocket ship race to claim a #36;10 million...