ChromeDBとBeluuuuuuga/Japanese-Instruction-Linux-Command-169を使って質問に対する適切なLinuxコマンドを探すRAGを構築する

初めに

珍しい?データセットがあったので、こちらを使ってLinuxで操作したいことがあるけどコマンドがわからない場合に雑に質問ができるRAGシステムを構築します

huggingface.co

実行した環境は以下でColobのNoteを公開しています

colab.research.google.com

デモ

開発環境

Chromeとは

Chromaはオープンソースかつローカルで使える埋め込みデータベースです。

docs.trychroma.com

RAGシステムの構築

ライブラリの準備

まずは以下でライブラリを入れます

%pip install -Uq chromadb numpy datasets

DBの構築とデータの追加

DBを作成します

# Import Chroma and instantiate a client. The default Chroma client is ephemeral, meaning it will not save to disk.
import chromadb

client = chromadb.Client()

# Create a new Chroma collection to store the supporting evidence. We don't need to specify an embedding fuction, and the default will be used.
collection = client.create_collection("Japanese-Instruction-Linux-Command")

huggingfaceからデータを取得して追加します

# Get the SciQ dataset from HuggingFace
from datasets import load_dataset

dataset = load_dataset("Beluuuuuuga/Japanese-Instruction-Linux-Command-169", split="train")


print("Number of questions with support: ", len(dataset))

# Embed and store the first 100 supports for this demo
i = 0
for data in dataset:
  collection.add(
      ids=[str(i)],  # IDs are just strings
      documents=data['instruction'],
      metadatas=[{"output": data['output']}],
  )
  i = i + 1

質問をして回答を取得

以下で質問に対して近い回答を取得します

results = collection.query(query_texts=dataset["instruction"][0], n_results=3)
print("質問:", dataset["instruction"][0])

# distanceのしきい値を設定
threshold = 0.05

# 検索結果をdistanceでフィルタリング
filtered_results = [result for result in zip(results['distances'][0], results['metadatas'][0]) if result[0] <= threshold]

if filtered_results:
    print("回答:")
    for distance, metadata in filtered_results:
        print(f"  distance: {distance:.4f}, output: {metadata['output']}")
else:
    print("適切な回答が見つかりませんでした。")

この場合の回答は以下になります

質問: Linuxで現在のディレクトリの内容を表示するコマンドを教えてください。
回答:
  distance: 0.0000, output: ls
  distance: 0.0360, output: pwd

ここで全く関係のないことを聞いてみます

q = "今日はいい天気ですね!"
results = collection.query(query_texts=q, n_results=3)
print("質問:", q)

# distanceのしきい値を設定
threshold = 0.05

# 検索結果をdistanceでフィルタリング
filtered_results = [result for result in zip(results['distances'][0], results['metadatas'][0]) if result[0] <= threshold]

if filtered_results:
    print("回答:")
    for distance, metadata in filtered_results:
        print(f"  distance: {distance:.4f}, output: {metadata['output']}")
else:
    print("適切な回答が見つかりませんでした。")

この場合は、見つからないので以下のようになります

質問: 今日はいい天気ですね!
適切な回答が見つかりませんでした。