SELECT 'a fat cat sat on a mat and ate a fat rat'::tsvector @@ 'cat & rat'::tsquery;
 ?column?
----------
 t

SELECT 'cat & rat'::tsquery @@ 'a fat cat sat on a mat and ate a fat rat'::tsvector;
 ?column?
----------
 t

CREATE TABLE items (
  id bigserial PRIMARY KEY,
  content text NOT NULL,
  embedding vector(3) NOT NULL -- 3 维
);

INSERT INTO items (content, embedding) VALUES
  ('a fat cat sat on a mat and ate a fat rat', '[1, 2, 3]'),
  ('a fat dog sat on a mat and ate a fat rat', '[4, 5, 6]'),
  ('a thin cat sat on a mat and ate a thin rat', '[7, 8, 9]'),
  ('a thin dog sat on a mat and ate a thin rat', '[10, 11, 12]');

SELECT content, embedding
FROM items
WHERE to_tsvector('english', content) @@ 'cat & rat'::tsquery;

SELECT content, embedding
FROM items
WHERE to_tsvector('english', content) @@ 'cat & rat'::tsquery
ORDER BY embedding <-> '[7, 8, 9]'::vector;

-- 查询从名为 "semantic_search" 的公用表表达式（CTE）开始。
WITH semantic_search AS (
    SELECT id, content, embedding, 
    RANK () OVER (ORDER BY embedding <-> '[7, 8, 9]') AS rank
    FROM items
    ORDER BY embedding <-> '[7, 8, 9]'::vector
    ), 
-- 定义另一个名为 "full_text_search" 的 CTE。
    full_text_search AS (
    SELECT id, content, embedding, 
    RANK () OVER (ORDER BY 
        ts_rank(to_tsvector('english', content), 
            'cat & rat'::tsquery) DESC) AS rank
    FROM items
    WHERE to_tsvector('english', content) @@ 'cat & rat'::tsquery
    ORDER BY ts_rank(
        to_tsvector('english', content), 'cat & rat'::tsquery) DESC
    )
-- 主查询从两个 CTE 中选择列并融合排名。
    SELECT 
        COALESCE(semantic_search.id, full_text_search.id) AS id,
        COALESCE(semantic_search.content, full_text_search.content) AS content,
        COALESCE(semantic_search.embedding, full_text_search.embedding) AS embedding,
        COALESCE(1.0 / (1 + semantic_search.rank), 0.0) + 
        COALESCE(1.0 / (1 + full_text_search.rank), 0.0) AS rank
    FROM semantic_search FULL OUTER JOIN full_text_search USING (id)
    ORDER BY rank DESC;

CREATE TABLE openai3072 (
  id bigserial PRIMARY KEY,
  text_embedding_3_large_3072_embedding vector(3072),
  text_embedding_3_large_3072_embedding bvector(3072)
);

-- 创建 256 维表达式索引
CREATE INDEX openai_vector_index_256 ON public.openai3072 USING vectors (((text_embedding_3_large_3072_embedding[0:256])::vector(256)) vector_l2_ops);

-- 创建 1024 维表达式索引
CREATE INDEX openai_vector_index_1024 on openai3072 using vectors((text_embedding_3_large_3072_embedding[0:1024]::vector(1024)) vector_l2_ops);

-- 创建 3072 维完整索引
CREATE INDEX openai_vector_index on openai3072 using vectors(text_embedding_3_large_3072_embedding vector_l2_ops);

CREATE INDEX openai_vector_index_bvector ON public.openai3072 USING vectors (text_embedding_3_large_3072_bvector bvector_l2_ops);

CREATE OR REPLACE FUNCTION match_documents_adaptive(
  query_embedding vector(3072),
  match_count int
)
RETURNS SETOF openai3072
LANGUAGE SQL
AS $$
WITH shortlist AS (
  SELECT *
  FROM openai3072
  ORDER BY (text_embedding_3_large_3072_embedding[0:256])::vector(256) <-> (query_embedding[0:256])::vector(256)
  LIMIT match_count * 2
)
SELECT *
FROM shortlist
ORDER BY text_embedding_3_large_3072_embedding <-> query_embedding
LIMIT match_count;
$$;

from transformers import CLIPProcessor, CLIPModel

processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")

def compute_image_embeddings(list_of_images):
    return model.get_image_features(**processor(images=list_of_images, return_tensors="pt", padding=True))

-- 创建表
CREATE TABLE images (
    id SERIAL PRIMARY KEY,
    author TEXT NOT NULL,
    vector vector(512) NOT NULL -- 512 维
);

INSERT INTO images (author, embedding) VALUES
  ('Allen', '[...]');

SELECT * FROM images ORDER BY embedding <-> '[...]' LIMIT 5;

SET vectors.hnsw_ef_search = 64;
SELECT * FROM images ORDER BY embedding <-> '[...]' WHERE author != 'Allen' LIMIT 64;

SELECT a.*, s.*, e.*
FROM asset_entity AS a
INNER JOIN smart_search AS s ON s.assetId = a.id
LEFT JOIN exif_info AS e ON e.assetId = a.id
WHERE a.ownerId IN (:...userIds)
  AND a.isArchived = false
  AND a.isVisible = true
  AND a.fileCreatedAt < NOW()
ORDER BY s.embedding <=> :embedding
LIMIT :numResults;

SELECT a.*, s.*, e.*

FROM asset_entity AS a
INNER JOIN smart_search AS s ON s.assetId = a.id
LEFT JOIN exif_info AS e ON e.assetId = a.id

WHERE a.ownerId IN (:...userIds)
  AND a.isArchived = false
  AND a.isVisible = true
  AND a.fileCreatedAt < NOW()

ORDER BY s.embedding <=> :embedding

LIMIT :numResults;

CREATE TABLE documents (
    id SERIAL PRIMARY KEY,
    user_id INTEGER NOT NULL,
    title TEXT,
    content TEXT,
    -- 其他元数据字段
);

CREATE TABLE chunks (
    id SERIAL PRIMARY KEY,
    document_id INTEGER REFERENCES documents(id),
    embedding vector(512) NOT NULL
);

CREATE INDEX idx_chunks_embedding ON chunks USING vectors (embedding vector_dot_ops);

SELECT chunks.id AS chunk_id, documents.id AS document_id, documents.title
FROM chunks
INNER JOIN documents ON chunks.document_id = documents.id
WHERE documents.user_id = 'alice' /* 替换为实际的 user_id */
ORDER BY chunks.embedding <-> '[3,2,1]' /* 替换为查询向量 */ LIMIT 5;

# pip install -U FlagEmbedding

from FlagEmbedding import BGEM3FlagModel

model = BGEM3FlagModel("BAAI/bge-m3", use_fp16=True)
sentences = ["What is BM25?", "Definition of BM25"]

output = model.encode(sentences, return_dense=True, return_sparse=True)

# 上述两段文本的稠密嵌入：
dense_embedding = output['dense_vecs']
# [[-0.02501617 -0.04525582 -0.01972856 ... -0.0099566   0.01677509
#   0.01292699]
# [-0.0106499  -0.04731942 -0.01477693 ... -0.0149255   0.01040097
#   0.00965083]

sparse_weight_id = output['lexical_weights']

# 您可以查看每个 token 对应的权重：
# print(model.convert_id_to_token(sparse_weight_id))
# 
# [{'What': 0.10430284, 'is': 0.10090457, 'BM': 0.2635918, '25': 0.3382988, '?': 0.052101523}, 
# {'Definition': 0.14298248, 'of': 0.076763615, 'BM': 0.2577639, '25': 0.33806682}]

# pip install -U numpy

import numpy as np

# 提取文本的稀疏向量下标与对应权重值
sparse_indices = [[int(token_id) for token_id in text] for text in sparse_weight_id]
# [[4865, 83, 90017, 2588, 32], [155455, 111, 90017, 2588]]
sparse_values = [[float(text[token_id]) for token_id in text] for text in sparse_weight_id]
# [[0.10430284, 0.10090457, 0.2635918, 0.3382988, 0.052101523], [0.14298248, 0.076763615, 0.2577639, 0.33806682]]

# 在 pgvecto.rs 中，每个稀疏向量的下标必须按升序排列
for i in range(len(sparse_weight_id)):
    sparse_values[i] = [x for _, x in sorted(zip(sparse_indices[i], sparse_values[i]))]
    sparse_indices[i] = sorted(sparse_indices[i])

docker run \
  --name pgvecto-rs-demo \
  -e POSTGRES_PASSWORD=mysecretpassword \
  -p 5432:5432 \
  -d tensorchord/pgvecto-rs:pg16-v0.3.0

print(len(model.tokenizer))
# 250002

# pip install -U psycopg

import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

with psycopg.connect(URL) as conn:
    conn.execute("DROP TABLE IF EXISTS documents;")
    conn.execute(f"""
        CREATE TABLE documents (
            id SERIAL PRIMARY KEY, 
            text TEXT NOT NULL,
            dense vector(1024) NOT NULL,
            sparse svector({len(model.tokenizer)}) NOT NULL);
    """
    )

# pip install -U pgvecto_rs

import psycopg
from pgvecto_rs.psycopg import register_vector
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

with psycopg.connect(URL) as conn:
    register_vector(conn)
    for text, dense, sparse_ind, sparse_val in zip(sentences, dense_embedding, sparse_indices, sparse_values):
        conn.execute(
            "INSERT INTO documents (text, dense, sparse) VALUES (%s, %s, to_svector(%s, %s::real[]::int[], %s::real[]));",
            (text, dense, len(model.tokenizer), sparse_ind, sparse_val),
        )
    conn.execute("""
        CREATE INDEX ON documents 
        USING vectors (sparse svector_dot_ops) 
        WITH (options = \"[indexing.hnsw]\");
    """)
    conn.execute("""
        CREATE INDEX ON documents 
        USING vectors (dense vector_l2_ops) 
        WITH (options = \"[indexing.hnsw]\");
    """)

def embedding(sentence: str) -> tuple[np.ndarray, list[list[int]], list[list[float]]]:
    """
    从文本生成稠密嵌入和稀疏嵌入
    """
    output = model.encode(sentence, return_dense=True, return_sparse=True, return_colbert_vecs=False)
    dense_embedding = output['dense_vecs']
    sparse_weight_id = output['lexical_weights']
    sparse_indices = [int(token_id) for token_id in sparse_weight_id]
    sparse_values = [float(sparse_weight_id[token_id]) for token_id in sparse_weight_id]

    sparse_values = [x for _, x in sorted(zip(sparse_indices, sparse_values))]
    sparse_indices = sorted(sparse_indices)
    return dense_embedding, sparse_indices, sparse_values

dense_embedding, sparse_indices, sparse_values = embedding("The text you want to search...")

with psycopg.connect(URL) as conn:
    register_vector(conn)
    cur = conn.execute(
        "SELECT text FROM documents ORDER BY sparse <#> to_svector(%s, %s::real[]::int[], %s::real[]) LIMIT 50;",
        (len(model.tokenizer), sparse_indices, sparse_values),
    )
    sparse_result = cur.fetchall()
      
    cur = conn.execute(
              "SELECT text FROM documents ORDER BY dense <-> %s LIMIT 50;",
              (dense_embedding,),
          )
    dense_result = cur.fetchall()
    
# 合并稀疏向量与稠密向量的搜索结果
mix_text = set([r[0] for r in dense_result]).union([r[0] for r in sparse_result])

from FlagEmbedding import FlagReranker

reranker = FlagReranker('BAAI/bge-reranker-large', use_fp16=True)

# 获取相关性得分（得分越高表示相关性越强）
scores = reranker.compute_score([["The text you want to search...", candidate] for candidate in mix_text])
# 对结合了稠密嵌入与稀疏嵌入检索出的候选文本进行重排序
reranked_text = [t for _, t in sorted(zip(scores, mix_text))]

# pip install openai

import openai

openai.api_key = "YOUR_API_KEY_HERE"
def complete(prompt):
    result = openai.Completion.create(
        engine='gpt-3.5-turbo-instruct',
        prompt=prompt,
        temperature=0,
        max_tokens=400,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0,
        stop=None
    )
    return result['choices'][0]['text'].strip()

base_prompt = "What is the release year and overall plot of Skyfall?"
print(complete(base_prompt))

base_prompt_2 = "What is the release year and overall plot of Cat Napping?"
print(complete(base_prompt_2))

# pip install -U datasets
from datasets import load_dataset
dataset = load_dataset("vishnupriyavr/wiki-movie-plots-with-summaries-faiss-embeddings", split='train')

texts = dataset['train']['text']
# texts[0]
# 'Alice in Wonderland \n  starring May Clark \n  released in the year1903 \n  with the following plot  \n Alice follows a large white rabbit down a "Rabbit-hole". She finds a tiny door. When she finds a bottle labeled "Drink me", she does, and shrinks, but not enough to pass through the door. She then eats something labeled "Eat me" and grows larger. She finds a fan when enables her to shrink enough to get into the "Garden" and try to get a "Dog" to play with her. She enters the "White Rabbit\'s tiny House," but suddenly resumes her normal size. In order to get out, she has to use the "magic fan."\r\nShe enters a kitchen, in which there is a cook and a woman holding a baby. She persuades the woman to give her the child and takes the infant outside after the cook starts throwing things around. The baby then turns into a pig and squirms out of her grip. "The Duchess\'s Cheshire Cat" appears and disappears a couple of times to Alice and directs her to the Mad Hatter\'s "Mad Tea-Party." After a while, she leaves.\r\nThe Queen invites Alice to join the "ROYAL PROCESSION": a parade of marching playing cards and others headed by the White Rabbit. When Alice "unintentionally offends the Queen", the latter summons the "Executioner". Alice "boxes the ears", then flees when all the playing cards come for her. Then she wakes up and realizes it was all a dream.'

embeddings = dataset['train']['embeddings']

# pip install cohere
import getpass
import os
import cohere

if not os.getenv("COHERE_API_KEY"):
    os.environ["COHERE_API_KEY"] = getpass.getpass("Enter your Cohere API key: ")

co = cohere.Client()
# 嵌入模型可参考此文档：https://docs.cohere.com/docs/cohere-embed
embed = co.embed(
    texts=texts,
    input_type="search_query",
    model="embed-english-v3.0"
)
embeddings = embed.embeddings

docker run \
  --name pgvecto-rs-demo \
  -e POSTGRES_PASSWORD=mysecretpassword \
  -p 5432:5432 \
  -d tensorchord/pgvecto-rs:pg16-v0.3.0

#pip install "psycopg[binary,pool]"
import psycopg

URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"
# 启用扩展
with psycopg.connect(URL) as conn:
    conn.execute("DROP EXTENSION IF EXISTS vectors;")
    conn.execute("CREATE EXTENSION vectors;")

import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

# 创建表
with psycopg.connect(URL) as conn:
    conn.execute("DROP TABLE IF EXISTS movies;")
    conn.execute(f"""
        CREATE TABLE movies (
            id SERIAL PRIMARY KEY, 
            text TEXT NOT NULL,
            embedding vector(768) NOT NULL);
    """
    )

# 可选安装，但推荐使用以显示进度条：
# pip install tqdm
import psycopg
from tqdm import tqdm
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

# 插入数据到表中
with psycopg.connect(URL) as conn:
    with conn.cursor() as cursor:    
        for text, embedding in tqdm(zip(texts, embeddings), total=len(texts)):
            cursor.execute(
                """INSERT INTO movies (text, embedding) 
                    VALUES (%s, %s::real[]);""",
                (text, embedding),
            )
        conn.commit()

import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

# 为 embedding 列创建索引
with psycopg.connect(URL) as conn:
    cursor = conn.cursor()
    cursor.execute("""
        CREATE INDEX ON movies 
        USING vectors (embedding vector_l2_ops) 
        WITH (options = \"[indexing.hnsw]\");
    """)
    conn.commit()

# pip install sentence-transformers

from sentence_transformers import SentenceTransformer

# 可选模型：'paraphrase-mpnet-base-v2', 'all-mpnet-base-v2'
encoder = SentenceTransformer('all-mpnet-base-v2')

def encode_query(query): # query 为字符串
    # 预处理查询语句
    query = query.lower()
    query = query.split()
    query = [word for word in query]
    query = ' '.join(query)
    
    # 对查询进行编码
    vector = encoder.encode(query)
    return vector

import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

def add_context(prompt):
    search_vector = encode_query(prompt).tolist()
    # 构造包含检索上下文的提示词
    prompt_start = (
        "Answer the question based on the following context:\n\n"
    )
    prompt_end = (
        f"\n\nQuestion: {prompt}\nAnswer:"
    )
    with psycopg.connect(URL) as conn:
        cursor = conn.cursor()
        # 根据希望提供给模型的电影文本数量调整 LIMIT 值
        cursor.execute("SELECT text FROM movies ORDER BY embedding <-> %s::real[] LIMIT 3;", 
                  (search_vector,),)
        result = cursor.fetchall()
    context = ""
    for r in result:
        context += r[0] + "\n\n"
    return prompt_start + context + prompt_end

base_prompt_2 = "What is the release year and overall plot of Cat Napping?"

context_prompt_2 = add_context(base_prompt_2)
context_response_2 = complete(context_prompt_2)
print(context_response_2)

base_prompt_3 = "In Cat Napping, what does Tom do to Jerry?"
print(complete(add_context(base_prompt_3)))