🌊

30Mの軽量モデルで文章をベクトル化し、質問に対して上位3つを返すものを作成。

に公開

ドキュメントを最初に計算しておく

import torch
import torch.nn.functional as F
from sentence_transformers import SentenceTransformer
import numpy as np
import csv

device = "cuda" if torch.cuda.is_available() else "cpu"
model = SentenceTransformer("cl-nagoya/ruri-v3-30m", device=device)

# ----------------------------
# 1. 事前にコーパスをベクトル化して保存
# ----------------------------
documents = [
    "川べりでサーフボードを持った人たちがいます",
    "サーファーたちが川べりに立っています",
    "瑠璃色(るりいろ)は、紫みを帯びた濃い青。名は、半貴石の瑠璃(ラピスラズリ)による。",
    "海は危険ですので気をつけましょう",
    "ミドリムシ大量発生中",
    "人気のゲームが発売中止??裏事情を徹底解説"
]

# 改行削除
documents = [doc.replace("\n", " ") for doc in documents]

# ベクトル化
embeddings = model.encode(documents, convert_to_tensor=False)

# 保存
np.savetxt("embeddings.csv", embeddings, fmt="%.6f", delimiter=",")

with open("texts.csv", "w", newline="", encoding="utf-8") as f:
    writer = csv.writer(f)
    for doc in documents:
        writer.writerow([doc])

print("保存完了: embeddings.csv, texts.csv")


# ----------------------------
# 2. 質問を与えて一番近い情報を検索(上位3件)
# ----------------------------
def search(query, top_k=3):
    # 質問のベクトル化
    query_emb = model.encode([query], convert_to_tensor=True)

    # 保存したベクトルをロード
    doc_embeddings = np.loadtxt("embeddings.csv", delimiter=",")
    doc_embeddings = torch.tensor(doc_embeddings, dtype=torch.float32)

    # 類似度計算
    similarities = F.cosine_similarity(query_emb, doc_embeddings)

    # 上位 top_k 件を取得
    topk = torch.topk(similarities, k=top_k)
    topk_idx = topk.indices.tolist()
    topk_scores = topk.values.tolist()

    # texts.csv から対応する文章を取得
    with open("texts.csv", "r", encoding="utf-8") as f:
        reader = list(csv.reader(f))
        results = [(reader[i][0], score) for i, score in zip(topk_idx, topk_scores)]

    return results


# ----------------------------
# テスト
# ----------------------------
query = "発売中止情報について教えて"
results = search(query, top_k=3)

print("質問:", query)
print("上位3件:")
for text, score in results:
    print(f"- {text} (類似度: {score:.4f})")

上記のコードを実行し、文書ベクトルを得て、
計算し終わったドキュメントと、質問文の内容をベクトル化したものをすべて計算し、
一番近いものを返す。

RAGの場合は、下のコードにGROQのコードとかを足して、質問を一回ベクトル化し、3件くらいとって、以上の情報をもとに以下の質問に答えてみたいな感じで実装していただければ行けると思います。


import torch
import torch.nn.functional as F
from sentence_transformers import SentenceTransformer
import numpy as np
import csv

device = "cuda" if torch.cuda.is_available() else "cpu"
model = SentenceTransformer("cl-nagoya/ruri-v3-30m", device=device)

# ----------------------------
# 保存したデータを使って検索する関数
# ----------------------------
def search(query, top_k=3, emb_file="embeddings.csv", txt_file="texts.csv"):
    # 質問のベクトル化
    query_emb = model.encode([query], convert_to_tensor=True)

    # 保存したベクトルをロード
    doc_embeddings = np.loadtxt(emb_file, delimiter=",")
    doc_embeddings = torch.tensor(doc_embeddings, dtype=torch.float32)

    # 類似度計算
    similarities = F.cosine_similarity(query_emb, doc_embeddings)

    # 上位 top_k 件を取得
    topk = torch.topk(similarities, k=top_k)
    topk_idx = topk.indices.tolist()
    topk_scores = topk.values.tolist()

    # texts.csv から対応する文章を取得
    with open(txt_file, "r", encoding="utf-8") as f:
        reader = list(csv.reader(f))
        results = [(reader[i][0], score) for i, score in zip(topk_idx, topk_scores)]

    return results


# ----------------------------
# 対話ループ
# ----------------------------
print("質問応答システム開始(終了するには 'exit' を入力してください)")
while True:
    query = input("\n質問> ")
    if query.lower() in ["exit", "quit", "終了"]:
        print("終了します。")
        break

    results = search(query, top_k=3)
    print("\n上位3件の関連情報:")
    for text, score in results:
        print(f"- {text} (類似度: {score:.4f})")


Discussion