RAG: Boosting LLMs with Contextual Retrieval

pip install --upgrade 'pymilvus[model]==2.4.4' 'numpy<2' openai requests
# pipenv install -v 'pymilvus[model]==2.4.4' 'numpy<2'  openai requests

from pymilvus.model.dense import SentenceTransformerEmbeddingFunction  # Sentence Transformer pre-trained models

# If connection to https://huggingface.co/ failed, uncomment the following path
# import os
# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

ef = SentenceTransformerEmbeddingFunction(
    model_name='all-MiniLM-L6-v2',  # Specify the model name
)

from pymilvus import MilvusClient, DataType

COLLECTION_NAME = "my_rag_collection"
SERVER_ADDR = "http://localhost:19530"
ACCESS_TOKEN = "root:Milvus"
DB_NAME = "default"

# 1. Set up a Milvus client
client = MilvusClient(
    uri=SERVER_ADDR,
    token=ACCESS_TOKEN,
    db_name=DB_NAME,
)

# 2. Check if the collection already exists and drop it if it does.
if client.has_collection(COLLECTION_NAME):
    client.drop_collection(COLLECTION_NAME)

# 3. Create a new collection with specified parameters.
client.create_collection(
    collection_name=COLLECTION_NAME,
    dimension=384,  # The vector has 384 dimensions, matching the SBERT embedding function with all-MiniLM-L6-v2
    auto_id=True,  # default is False
    # primary_field_name="id",
    # id_type="int",
    # vector_field_name="vector",
    # metric_type="COSINE",
    # enable_dynamic_field=True,
)

# 4. (optional) To load a collection, use the load_collection() method.
# client.load_collection(
#     collection_name=COLLECTION_NAME
# )
#
# To release a collection, use the release_collection() method.
# client.release_collection(
#     collection_name=COLLECTION_NAME
# )

# 5. (optional) The collection created above is loaded automatically.
res = client.get_load_state(
    collection_name=COLLECTION_NAME
)

print(res)

# 6. (optional) List detailed information about the collection.
import json
desc = client.describe_collection(
    collection_name=COLLECTION_NAME,
)
print(json.dumps(desc, indent=2))

{'state': <LoadState: Loaded>}
{
  "collection_name": "my_rag_collection",
  "auto_id": true,
  "num_shards": 1,
  "description": "",
  "fields": [
    {
      "field_id": 100,
      "name": "id",
      "description": "",
      "type": 5,
      "params": {},
      "auto_id": true,
      "is_primary": true
    },
    {
      "field_id": 101,
      "name": "vector",
      "description": "",
      "type": 101,
      "params": {
        "dim": 384
      }
    }
  ],
  "aliases": [],
  "collection_id": 450568843972908135,
  "consistency_level": 2,
  "properties": {},
  "num_partitions": 1,
  "enable_dynamic_field": true
}

# download and save it as a local text file.
import os
import urllib.request

URL = "https://raw.githubusercontent.com/milvus-io/milvus/master/DEVELOPMENT.md"
FILE_PATH = "./Milvus_DEVELOPMENT.md1"

if not os.path.exists(FILE_PATH):
    urllib.request.urlretrieve(URL, FILE_PATH)

from pymilvus import MilvusClient, model

COLLECTION_NAME = "my_rag_collection"
SERVER_ADDR = "http://localhost:19530"
ACCESS_TOKEN = "root:Milvus"
DB_NAME = "default"

client = MilvusClient(
    uri=SERVER_ADDR,
    token=ACCESS_TOKEN,
    db_name=DB_NAME,
)

# If connection to https://huggingface.co/ failed, uncomment the following path
# import os
# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

ef = model.dense.SentenceTransformerEmbeddingFunction(
    model_name='all-MiniLM-L6-v2',
)

FILE_PATH = "./Milvus_DEVELOPMENT.md"

with open(FILE_PATH, "r+t", encoding='utf-8') as fi:
    text = fi.read()

# Use "# " to separate the content in the file, which can roughly separate
# the content of each main part of the markdown file.
docs = text.split("# ")

vectors = ef.encode_documents(docs)
data = [{"vector": vectors[i], "text": docs[i]} for i in range(len(vectors))]

res = client.insert(collection_name=COLLECTION_NAME, data=data)
print(res)

{'insert_count': 47, 'ids': [450568843971283844, ... , 450568843971283889, 450568843971283890], 'cost': 0}