开发基于云的RAG应用，使用开源 LLM

pip install bentoml langchain clickhouse-connect

from langchain_community.document_loaders.wikipedia import WikipediaLoader
loader = WikipediaLoader(query="Albert Einstein")

# Load the documents
docs = loader.load()

# Display the content of the first document
print(docs[0].page_content)

from langchain_text_splitters import CharacterTextSplitter
# Split the text into chunks
text = ' '.join([page.page_content.replace('\\t', ' ') for page in docs])
text_splitter = CharacterTextSplitter(
   separator="\\n",
   chunk_size=400,
   chunk_overlap=100,
   length_function=len,
   is_separator_regex=False,
)
texts = text_splitter.create_documents([text])
splits = [item.page_content for item in texts]

# Import the libraries
import subprocess
import sys
import numpy as np

# Install the packages if the API key isn't provided
def install(package):
   subprocess.check_call([sys.executable, "-m", "pip", "install", package])

# Define the embedding method
def get_embeddings(texts: list, BENTO_EMBEDDING_MODEL_END_POINT=None, BENTO_API_TOKEN=None) -> list:
   # If the BentoML KEY is provided, the method will use BENTOML model to get embeddings
   if BENTO_EMBEDDING_MODEL_END_POINT and BENTO_API_TOKEN:
       import bentoml
       embedding_client = bentoml.SyncHTTPClient(BENTO_EMBEDDING_MODEL_END_POINT, token=BENTO_API_TOKEN)
       return embedding_client.encode(sentences=texts).tolist()
  
   # Otherwise it'll use transformers library
   else:
       # Install transformers and torch if not already installed
       try:
           import transformers
       except ImportError:
           install("transformers")
       try:
           import torch
       except ImportError:
           install("torch")
      
       from transformers import AutoTokenizer, AutoModel
      
       # Initialize the tokenizer and model for embeddings
       tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
       model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
      
       inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt", max_length=512)
       with torch.no_grad():
           outputs = model(**inputs)
       embeddings = outputs.last_hidden_state.mean(dim=1)
       return embeddings.numpy().tolist()

all_embeddings = []
# Pass the splits in a batch of 25
for i in range(0, len(splits), 25):
   batch = splits[i:i+25]
   # Pass the batch to the get_embeddings method
   embeddings_batch = get_embeddings(batch)
   # Append the embeddings to the all_embeddings list holdng the embeddings of the whole dataset
   all_embeddings.extend(embeddings_batch)

import pandas as pd
df = pd.DataFrame({
   'page_content': splits,
   'embeddings': all_embeddings
})

import clickhouse_connect
client = clickhouse_connect.get_client(
   host='your-host-name',
   port=443,
   username='your-user-name',
   password='your-password'
)

# Create the table named RAG
client.command("""
CREATE TABLE IF NOT EXISTS default.RAG (
   id Int64,
   page_content String,
   embeddings Array(Float32),
   CONSTRAINT check_data_length CHECK length(embeddings) = 384
) ENGINE = MergeTree()
   ORDER BY id
""")

# Insert data into the table
batch_size = 100
num_batches = (len(df) + batch_size - 1) // batch_size

for i in range(num_batches):
   batch_data = df[i * batch_size: (i + 1) * batch_size]
   client.insert('default.RAG', batch_data.values.tolist(), column_names=batch_data.columns.tolist())
   print(f"Batch {i+1}/{num_batches} inserted.")

client.command("""
ALTER TABLE default.RAG
   ADD VECTOR INDEX vector_index embeddings
   TYPE MSTG
""")

def get_relevant_docs(user_query, top_k):
   query_embeddings = get_embeddings(user_query)[0]
   results = client.query(f"""
       SELECT page_content,
       distance(embeddings, {query_embeddings}) as dist FROM default.RAG ORDER BY dist LIMIT {top_k}
   """)
   relevant_docs = " "
  
   for row in results.named_results():
       relevant_docs=relevant_docs + row["page_content"]
  
   return relevant_docs

# Example query
message="Who is albert einstein?"
relevant_docs = get_relevant_docs(message, 8)
print(relevant_docs)

import bentoml
BENTO_LLM_END_POINT = "add-your-end-point-here"

llm_client = bentoml.SyncHTTPClient(BENTO_LLM_END_POINT, token="your-token-here")

def dorag(question: str, context: str):
   # Define the prompt template
   prompt = (f"You are a helpful assistant. The user has a question. Answer the user question based only on the context: {context}. \\n"
             f"The user question is {question}")
   # Call the LLM endpoint with the prompt defined above
   results = llm_client.generate(
       max_tokens=1024,
       prompt=prompt,
   )
   res = ""
   for result in results:
       res += result
   return res

query = "Who is albert einstein?"
dorag(question=query, context=relevant_docs)