以自组织映射算法促进高效的LLM检索增强生成

class KohonenSOM():
"""
该代码是基于以下文章开发的：
http://www.ai-junkie.com/ann/som/som1.html

向量和矩阵运算是使用PyTorch张量进行的。
"""
def __init__( ... )
...
def find_topk_best_matching_units( self, data_points : torch.Tensor, topk : int = 1 ) -> List[ List[ int ] ] :
if len( data_points.size() ) == 1:
#batching 
data_points = data_points.view( 1, data_points.shape[0] )

topk = int( topk )

distances = self.dist_evaluator( data_points, self.lattice_node_weights )

topk_best_matching_unit_indexes = torch.topk( distances, topk, dim=1, largest=False ).indices
topk_best_matching_units = []

for i in range( data_points.shape[0] ):
best_matching_unit_indexes = topk_best_matching_unit_indexes[i]
best_matching_units = [ self.lattice_coordinates[ bmu_index.item() ].tolist() for bmu_index in best_matching_unit_indexes ]
topk_best_matching_units.append( best_matching_units )

return topk_best_matching_units

class SOMBasedVectorIndexer():
...

def train_n_gen_indexes( 
self, input_vectors : torch.Tensor, 
train_epochs : int = 100 
):
if self.generated_indexes:
print( "WARNING: Indexes were already generated. Ignoring the request..." )
return

self.som.train( input_vectors, train_epochs )

topk_bmu_indexes = self.som.find_topk_best_matching_units( input_vectors, topk = self.topk_bmu_for_indexing )

for idx in tqdm( range( len( topk_bmu_indexes ) ), desc="SOM-Based Indexed Vectors"  ):
bmu_indexes = topk_bmu_indexes[ idx ]

for bmu_index in bmu_indexes:
bmu_index_key = tuple( bmu_index )
idx_set = self.som_node_idx_map.get( bmu_index_key, set() )
idx_set.add( idx )
self.som_node_idx_map[ bmu_index_key ] = idx_set

self.generated_indexes = True

import openai
from openai.embeddings_utils import get_embedding
...
from vector_encoder_parent import VectorEncoder
...

class OpenAIEmbeddingsVectorEncoder( VectorEncoder ):
def __init__( ... )
...
def encode_batch( self, list_of_text : List[ str ] ) -> torch.Tensor :
if list_of_text == None or len( list_of_text ) == 0:
raise ValueError( "ERROR: Required list_of_text is None or empty" )

list_of_text = [ str( text ) for text in list_of_text ]

openai.api_key = self.openai_key
response = openai.Embedding.create(
input = list_of_text,
engine = self.vector_encoder_id
)

embeddings = [ data["embedding"] for data in response["data"] ] 
vectors = torch.tensor( embeddings, dtype=torch.float )
return vectors

import requests
import pandas as pd
from dateutil.parser import parse
...
class WikiEventsDataSource():
...
def fetch_n_prepare_data( self ):
if self.fetched:
print( "WARNING: Wiki events for the specified years already fetched. Ignoring the request..." )
return

main_df = pd.DataFrame()

for year in self.event_years_to_fetch:
wiki_api_params = {
"action": "query", 
"prop": "extracts",
"exlimit": 1,
"titles": year,
"explaintext": 1,
"formatversion": 2,
"format": "json"
}

response = requests.get( "https://en.wikipedia.org/w/api.php", params=wiki_api_params )
response_dict = response.json()

df = pd.DataFrame()
df[ "text" ] = response_dict["query"]["pages"][0]["extract"].split("\n")
df = self.__clean_df__( df, year )

main_df = pd.concat( [ main_df, df ] )

self.df = main_df.reset_index(drop=True)
self.fetched = True

...
from vector_encoder_parent import VectorEncoder
from vector_indexer import SOMBasedVectorIndexer

class SOM_Based_RAG_Util():
...
def load_n_vectorize_data( self, data_source ):
if self.data_loaded_n_vectorized:
print( "WARNING: Data already loaded and vectorized. Ignoring the request..." )
return

data_source.fetch_n_prepare_data()
self.df = data_source.get_data()

vectors = None

for i in tqdm( range(0, len(self.df), self.vectorize_batch_size ), desc="Vectorized Data Batch" ):
list_of_text = self.df.iloc[ i:i+self.vectorize_batch_size ]["text"].tolist()
batch_encoded_vectors = self.vector_encoder.encode_batch( list_of_text )

if vectors == None:
vectors = batch_encoded_vectors
else:
vectors = torch.cat( [ vectors, batch_encoded_vectors], dim=0 )

self.vectors = vectors.to( self.device )
self.data_loaded_n_vectorized = True

def train_n_index_data_vectors( self, train_epochs : int = 100  ):
if not self.data_loaded_n_vectorized:
raise ValueError( "ERROR: Data not loaded and vectorized." )

if self.data_vectors_indexed:
print( "WARNING: Data vectors already indexed. Ignoring the request..." )
return

self.vector_indexer.train_n_gen_indexes( self.vectors, train_epochs )
self.data_vectors_indexed = True

def find_semantically_similar_data( self, query: str, sim_evaluator = None, sim_threshold : float = 0.8  ):
if not self.data_vectors_indexed:
raise ValueError( "ERROR: Data vectors not indexed." )

if query == None or len( query.strip() ) == 0:
raise ValueError( "ERROR: Required query text is not specified." )

sim_threshold = float( sim_threshold )

if sim_evaluator == None:
sim_evaluator = nn.CosineSimilarity(dim=0, eps=1e-6)

query_vector = self.vector_encoder.encode( query )
query_vector = query_vector.view( self.vector_encoder.get_encoded_vector_dimensions() )
query_vector = query_vector.to( self.device )

nearest_indexes = self.vector_indexer.find_nearest_indexes( query_vector )
nearest_indexes = nearest_indexes[0]

sim_scores = []

for idx in nearest_indexes:
data_vector = self.vectors[ idx ]
data_vector = data_vector.view( self.vector_encoder.get_encoded_vector_dimensions() )

sim_score = sim_evaluator( query_vector, data_vector )

if sim_score >= sim_threshold:
sim_score_tuple = (idx, sim_score.item() )
sim_scores.append( sim_score_tuple )

sim_scores.sort( key = lambda x: x[1], reverse=True )

semantically_similar_data = [ 
{ 
'text': self.df[ 'text' ][ idx ],
'sim_score' : sim_score
} for idx, sim_score in sim_scores
]

return semantically_similar_data

from abc import ABC, abstractmethod
import torch
import math

class QuestionAnswerChatBot( ABC ):
...
def find_answer_to_question( self, question : str, sim_threshold = 0.68, max_new_tokens : int = 5 ):
if question == None or len( question.strip() ) == 0:
raise ValueError( "ERROR: Required question is not specified" )

sim_threshold = float( sim_threshold )
max_new_tokens = int( max_new_tokens )

qa_instruction = self.get_qa_instruction( question, sim_threshold = sim_threshold )

answer_text = self.__get_answer_text__( qa_instruction, max_new_tokens = max_new_tokens )
answer_text = self.__clean_answer_text__( qa_instruction, answer_text )

return answer_text
...
def __qa_template__( self ):
qa_template = """Context: 

{}

---

Question: {}
Answer:"""
return qa_template

import openai
import tiktoken
from qa_chatbot import QuestionAnswerChatBot

class OpenAIQuestionAnswerChatBot( QuestionAnswerChatBot ):
...
def __get_answer_text__( self, qa_instruction : str, max_new_tokens : int = 5 ) -> str :
openai.api_key = self.openai_key

basic_answer = openai.Completion.create(
model = self.openai_model_name,
prompt = qa_instruction, 

)

answer_text = basic_answer[ "choices" ][0][ "text" ]
return answer_text

def __token_count__( self, text : str ):    
return len( self.tokenizer.encode( text ) )

sample_questions = [
"Who won the 2022 soccer world cup?",
"When did Sweden join NATO?",
"Who joined NATO in 2023?",
"Who joined NATO in 2024?",
"Which is the 31st member of NATO?",
"Which is the 32nd member of NATO?",
"Who won the Cricket World Cup in 2023?",
"Who defeated India in Cricket World Cup final in 2023?",
"Name the former prime minister of Japan that was assassinated in 2022?",
"When did Chandrayaan-3 land near the south pole of the Moon?",
"Where did Chandrayaan-3 land on the Moon?",
"Who acquired Twitter in 2022?",
"Who owns Twitter?",
"Who acquired Activision Blizzard in 2023?"
]

openai_vector_encoder_id = "text-embedding-ada-002"
openai_encoded_vector_dimensions = 1536
openai_tokenizer_name = "cl100k_base" 
openai_model_name = "gpt-3.5-turbo-instruct"

vector_encoder = OpenAIEmbeddingsVectorEncoder( openai_encoded_vector_dimensions, openai_vector_encoder_id, openai_key )

event_years_to_fetch = [ 2022, 2023, 2024 ]
data_source = WikiEventsDataSource( event_years_to_fetch  )
...
som_driven_rag_util = SOM_Based_RAG_Util( 
vector_encoder = vector_encoder,
som_lattice_height = 20,
som_lattice_width = 30,
learning_rate = 0.3,
topk_bmu_for_indexing = 10,
device = device
)
...
openai_chatbot = OpenAIQuestionAnswerChatBot( 
vector_db_util = som_driven_rag_util,
openai_tokenizer_name = openai_tokenizer_name,
openai_model_name = openai_model_name,
openai_key = openai_key,
question_input_max_token_count = 100,
context_trim_percent = 0.1,
device = device
)