chromadb example

chroma_sample.py

import os
import chromadb
chroma_client = chromadb.Client()

collection = chroma_client.create_collection(name="wiki")

def load_txt_files(directory):
    txt_files = []
    names = []

    for filename in os.listdir(directory):
        if filename.endswith(".txt"):
            names.append(filename)
            file_path = os.path.join(directory, filename)
            with open(file_path, 'r') as file:
                txt_files.append(file.read())

    return (txt_files, names)

# Replace 'directory_path' with the path of the directory containing your .txt files
docs, ids = load_txt_files('./wiki') # todo - move into a config file

collection.add(
    documents=docs,
    ids=ids
)

# First let's look for an exact match
results = collection.query(
    query_texts = ['geckos are usually nocturnal'],
    n_results = 1
)
print(results['ids'])


results = collection.query(
    query_texts = ['geckos are nocturnal'],
    n_results = 1
)
print(results['ids'])


results = collection.query(
    query_texts = ['geckos are nighttime'],
    n_results = 1
)

print(results['ids'])
results = collection.query(
    query_texts = ['geckos prefer nighttime'],
    n_results = 1
)
print(results['ids'])


print(results['ids'])
results = collection.query(
    query_texts = ['lizards prefer nighttime'],
    n_results = 1
)
print(results['ids'])