如何在生成式 AI 里使用 Ray Data 进行大规模 RAG 应用的 Embedding Inference

pip install boto3 pip install awscli –force-reinstall –upgradeaws config #fill in your key and secret keymore ~/.aws/credentials #make sure this looks correct

pip install -U "ray[data]"

pip install -U pymilvus "pymilvus[model]" langchain

import ray, os, pprint, time, boto3from langchain.text_splitter import RecursiveCharacterTextSplitterimport numpy as npimport pymilvusprint(pymilvus.__version__) # must be >= 2.4.0from pymilvus.model.hybrid import BGEM3EmbeddingFunction

chunk_size = 512chunk_overlap = np.round(chunk_size * 0.10, 0)
# Define a LangChain text splitter.text_splitter = RecursiveCharacterTextSplitter(   chunk_size=chunk_size,   chunk_overlap=chunk_overlap,   length_function=len)  #len is a built-in Python function
# 1. Define a regular python function for chunking.def chunk_row(row, splitter=text_splitter):
   # Copy the row columns into metadata.   metadata = row.copy()   del metadata['text'] # Remove text from metadata
   # Split the text into chunks.   chunks = splitter.create_documents(       texts=[row["text"]],       metadatas=[metadata])   chunk_list = [{       "text": chunk.page_content,       **chunk.metadata} for chunk in chunks]
   return chunk_list
# 2. Define a class with a callable method to compute embeddings.class ComputeEmbeddings:     def __init__(self):           # Initialize a Milvus built-in sparse-dense-late-interaction-reranking encoder.           # https://huggingface.co/BAAI/bge-m3           self.model = BGEM3EmbeddingFunction(use_fp16=False, device="cpu")           print(f"dense_dim: {self.model.dim['dense']}")           print(f"sparse_dim: {self.model.dim['sparse']}")
     def __call__(self, batch):
           # Ray data batch is a dictionary where values are array values.           # BGEM3EmbeddingFunction input is docs as a list of strings.           docs = list(batch['text'])
           # Encode the documents. bge-m3 dense embeddings are already normalized.           embeddings = self.model(docs)           batch['vector_dense'] = embeddings['dense']           return batch
if __name__ == "__main__":
   FILE_PATH = "s3://zilliz/kaggle_imdb.parquet"
   # Load and transform data.
   ds = ray.data.read_parquet(FILE_PATH)
   # Chunk the input text   chunked_ds = ds.flat_map(chunk_row)
   # Compute embeddings with a class that calls the embeddings model.   embeddings_ds = chunked_ds.map_batches(ComputeEmbeddings, concurrency=4)
   # Save the embeddings to S3 in a folder of parquet part files.   embeddings_ds.write_parquet('s3://zilliz/kaggle_imdb_embeddings')

def mc_run_search(question, output_fields, top_k=2, filter_expression=""):     # Embed the question using the same encoder.   embeddings = model_bgem3([question])   query_embeddings = embeddings['dense']
   # Run semantic vector search using your query and the vector database.   results = mc.search(       COLLECTION_NAME,       data=query_embeddings,       search_params=SEARCH_PARAMS,       output_fields=output_fields,       # Milvus can utilize metadata in boolean expressions to filter search.       filter=filter_expression,       limit=top_k,       consistency_level="Eventually"   )
   # Assemble retrieved context and context metadata.   # The search result is in the variable `results[0]`, which is type   # 'pymilvus.orm.search.SearchResult'.   METADATA_FIELDS = [f for f in output_fields if f != 'chunk']   formatted_results, context, context_metadata = _utils.client_assemble_retrieved_context(       results, metadata_fields=METADATA_FIELDS, num_shot_answers=top_k)   return formatted_results, context, context_metadata
SAMPLE_QUESTION = "muybridge horse movie"
# Return top k unique results with HNSW index.TOP_K = 2
# Define output fields to return.OUTPUT_FIELDS = ["movie_id", "chunk", "PosterLink"]
formatted_results, context, context_metadata = \   mc_run_search(SAMPLE_QUESTION, OUTPUT_FIELDS, TOP_K)