使用 Milvus、vLLM 和 Llama 3.1 搭建 RAG 应用

    from langchain.document_loaders import DirectoryLoader    # Load HTML files already saved in a local directory    path = "../../RAG/rtdocs_new/"    global_pattern = '*.html'    loader = DirectoryLoader(path=path, glob=global_pattern)    docs = loader.load()

    # Print num documents and a preview.    print(f"loaded {len(docs)} documents")    print(docs[0].page_content)    pprint.pprint(docs[0].metadata)

    import torch    from sentence_transformers import SentenceTransformer

    # Initialize torch settings for device-agnostic code.    N_GPU = torch.cuda.device_count()    DEVICE = torch.device('cuda:N_GPU' if torch.cuda.is_available() else 'cpu')

    # Download the model from huggingface model hub.    model_name = "BAAI/bge-large-en-v1.5"    encoder = SentenceTransformer(model_name, device=DEVICE)

    # Get the model parameters and save for later.    EMBEDDING_DIM = encoder.get_sentence_embedding_dimension()    MAX_SEQ_LENGTH_IN_TOKENS = encoder.get_max_seq_length()

    # Inspect model parameters.    print(f"model_name: {model_name}")    print(f"EMBEDDING_DIM: {EMBEDDING_DIM}")    print(f"MAX_SEQ_LENGTH: {MAX_SEQ_LENGTH}")

    from langchain.text_splitter import RecursiveCharacterTextSplitter


    CHUNK_SIZE = 512    chunk_overlap = np.round(CHUNK_SIZE * 0.10, 0)    print(f"chunk_size: {CHUNK_SIZE}, chunk_overlap: {chunk_overlap}")

    # Define the splitter.    child_splitter = RecursiveCharacterTextSplitter(       chunk_size=CHUNK_SIZE,       chunk_overlap=chunk_overlap)

    # Chunk the docs.    chunks = child_splitter.split_documents(docs)    print(f"{len(docs)} docs split into {len(chunks)} child documents.")

    # Encoder input is doc.page_content as strings.    list_of_strings = [doc.page_content for doc in chunks if hasattr(doc, 'page_content')]

    # Embedding inference using HuggingFace encoder.    embeddings = torch.tensor(encoder.encode(list_of_strings))

    # Normalize the embeddings.    embeddings = np.array(embeddings / np.linalg.norm(embeddings))

    # Milvus expects a list of `numpy.ndarray` of `numpy.float32` numbers.    converted_values = list(map(np.float32, embeddings))

    # Create dict_list for Milvus insertion.    dict_list = []    for chunk, vector in zip(chunks, converted_values):       # Assemble embedding vector, original text chunk, metadata.       chunk_dict = {           'chunk': chunk.page_content,           'source': chunk.metadata.get('source', ""),           'vector': vector,       }       dict_list.append(chunk_dict)

    # Connect a client to the Milvus Lite server.    from pymilvus import MilvusClient    mc = MilvusClient("milvus_demo.db")


    # Create a collection with flexible schema and AUTOINDEX.    COLLECTION_NAME = "MilvusDocs"    mc.create_collection(COLLECTION_NAME,           EMBEDDING_DIM,           consistency_level="Eventually",           auto_id=True,            overwrite=True)

    # Insert data into the Milvus collection.    print("Start inserting entities")    start_time = time.time()    mc.insert(       COLLECTION_NAME,       data=dict_list,       progress_bar=True)

    end_time = time.time()    print(f"Milvus insert time for {len(dict_list)} vectors: ", end="")    print(f"{round(end_time - start_time, 2)} seconds")

    SAMPLE_QUESTION = "What do the parameters for HNSW mean?"


    # Embed the question using the same encoder.    query_embeddings = torch.tensor(encoder.encode(SAMPLE_QUESTION))    # Normalize embeddings to unit length.    query_embeddings = F.normalize(query_embeddings, p=2, dim=1)    # Convert the embeddings to list of list of np.float32.    query_embeddings = list(map(np.float32, query_embeddings))

    # Define metadata fields you can filter on.    OUTPUT_FIELDS = list(dict_list[0].keys())    OUTPUT_FIELDS.remove('vector')

    # Define how many top-k results you want to retrieve.    TOP_K = 2

    # Run semantic vector search using your query and the vector database.    results = mc.search(        COLLECTION_NAME,        data=query_embeddings,        output_fields=OUTPUT_FIELDS,        limit=TOP_K,        consistency_level="Eventually")

    # (Recommended) Create a new conda environment.    conda create -n myenv python=3.11 -y    conda activate myenv

    # Install vLLM with CUDA 12.1.    pip install -U vllm transformers torch

    import vllm, torch    from vllm import LLM, SamplingParams    # Clear the GPU memory cache.    torch.cuda.empty_cache()    # Check the GPU.    !nvidia-smi

    # Login to HuggingFace using your new token.    from huggingface_hub import login    from google.colab import userdata    hf_token = userdata.get('HF_TOKEN')    login(token = hf_token, add_to_git_credential=True)

    # 1. Choose a model    MODELTORUN = "meta-llama/Meta-Llama-3.1-8B-Instruct"
    # 2. Clear the GPU memory cache, you're going to need it all!    torch.cuda.empty_cache()

    # 3. Instantiate a vLLM model instance.    llm = LLM(model=MODELTORUN,             enforce_eager=True,             dtype=torch.bfloat16,             gpu_memory_utilization=0.5,             max_model_len=1000,             seed=415,             max_num_batched_tokens=3000)

    # Separate all the context together by space.    contexts_combined = ' '.join(contexts)    # Lance Martin, LangChain, says put the best contexts at the end.    contexts_combined = ' '.join(reversed(contexts))

    # Separate all the unique sources together by comma.    source_combined = ' '.join(reversed(list(dict.fromkeys(sources))))

    SYSTEM_PROMPT = f"""First, check if the provided Context is relevant to    the user's question.  Second, only if the provided Context is strongly relevant, answer the question using the Context.  Otherwise, if the Context is not strongly relevant, answer the question without using the Context.     Be clear, concise, relevant.  Answer clearly, in fewer than 2 sentences.    Grounding sources: {source_combined}    Context: {contexts_combined}    User's question: {SAMPLE_QUESTION}    """

    prompts = [SYSTEM_PROMPT]

    # Sampling parameters    sampling_params = SamplingParams(temperature=0.2, top_p=0.95)

    # Invoke the vLLM model.    outputs = llm.generate(prompts, sampling_params)

    # Print the outputs.    for output in outputs:       prompt = output.prompt       generated_text = output.outputs[0].text       # !r calls repr(), which prints a string inside quotes.       print()       print(f"Question: {SAMPLE_QUESTION!r}")       pprint.pprint(f"Generated text: {generated_text!r}")