1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212

import fitz
import os
import numpy as np
import json
from openai import OpenAI

def extract_text_from_pdf(pdf_path):
    """
    Extracts text from a PDF file.

    Args:
    pdf_path (str): Path to the PDF file.

    Returns:
    str: Extracted text from the PDF.
    """
    # Open the PDF file
    mypdf = fitz.open(pdf_path)
    all_text = ""  # Initialize an empty string to store the extracted text

    # Iterate through each page in the PDF
    for page in mypdf:
        # Extract text from the current page and add spacing
        all_text += page.get_text("text") + " "

    # Return the extracted text, stripped of leading/trailing whitespace
    return all_text.strip()

# Initialize the OpenAI client with the base URL and API key
client = OpenAI(
    base_url="https://api.studio.nebius.com/v1/",
    api_key=os.getenv("OPENAI_API_KEY")  # Retrieve the API key from environment variables
)

def get_embedding(text, model="BAAI/bge-en-icl"):
    """
    Creates an embedding for the given text using OpenAI.

    Args:
    text (str): Input text.
    model (str): Embedding model name.

    Returns:
    np.ndarray: The embedding vector.
    """
    response = client.embeddings.create(model=model, input=text)
    return np.array(response.data[0].embedding)

def cosine_similarity(vec1, vec2):
    """
    Computes cosine similarity between two vectors.

    Args:
    vec1 (np.ndarray): First vector.
    vec2 (np.ndarray): Second vector.

    Returns:
    float: Cosine similarity.
    """
    return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))

def compute_breakpoints(similarities, method="percentile", threshold=90):
    """
    Computes chunking breakpoints based on similarity drops.

    Args:
    similarities (List[float]): List of similarity scores between sentences.
    method (str): 'percentile', 'standard_deviation', or 'interquartile'.
    threshold (float): Threshold value (percentile for 'percentile', std devs for 'standard_deviation').

    Returns:
    List[int]: Indices where chunk splits should occur.
    """
    # Determine the threshold value based on the selected method
    if method == "percentile":
        # Calculate the Xth percentile of the similarity scores
        threshold_value = np.percentile(similarities, threshold)
    elif method == "standard_deviation":
        # Calculate the mean and standard deviation of the similarity scores
        mean = np.mean(similarities)
        std_dev = np.std(similarities)
        # Set the threshold value to mean minus X standard deviations
        threshold_value = mean - (threshold * std_dev)
    elif method == "interquartile":
        # Calculate the first and third quartiles (Q1 and Q3)
        q1, q3 = np.percentile(similarities, [25, 75])
        # Set the threshold value using the IQR rule for outliers
        threshold_value = q1 - 1.5 * (q3 - q1)
    else:
        # Raise an error if an invalid method is provided
        raise ValueError("Invalid method. Choose 'percentile', 'standard_deviation', or 'interquartile'.")

    # Identify indices where similarity drops below the threshold value
    return [i for i, sim in enumerate(similarities) if sim < threshold_value]

def split_into_chunks(sentences, breakpoints):
    """
    Splits sentences into semantic chunks.

    Args:
    sentences (List[str]): List of sentences.
    breakpoints (List[int]): Indices where chunking should occur.

    Returns:
    List[str]: List of text chunks.
    """
    chunks = []  # Initialize an empty list to store the chunks
    start = 0  # Initialize the start index

    # Iterate through each breakpoint to create chunks
    for bp in breakpoints:
        # Append the chunk of sentences from start to the current breakpoint
        chunks.append(". ".join(sentences[start:bp + 1]) + ".")
        start = bp + 1  # Update the start index to the next sentence after the breakpoint

    # Append the remaining sentences as the last chunk
    chunks.append(". ".join(sentences[start:]))
    return chunks  # Return the list of chunks

def create_embeddings(text_chunks):
    """
    Creates embeddings for each text chunk.

    Args:
    text_chunks (List[str]): List of text chunks.

    Returns:
    List[np.ndarray]: List of embedding vectors.
    """
    # Generate embeddings for each text chunk using the get_embedding function
    return [get_embedding(chunk) for chunk in text_chunks]

def semantic_search(query, text_chunks, chunk_embeddings, k=5):
    """
    Finds the most relevant text chunks for a query.

    Args:
    query (str): Search query.
    text_chunks (List[str]): List of text chunks.
    chunk_embeddings (List[np.ndarray]): List of chunk embeddings.
    k (int): Number of top results to return.

    Returns:
    List[str]: Top-k relevant chunks.
    """
    # Generate an embedding for the query
    query_embedding = get_embedding(query)

    # Calculate cosine similarity between the query embedding and each chunk embedding
    similarities = [cosine_similarity(query_embedding, emb) for emb in chunk_embeddings]

    # Get the indices of the top-k most similar chunks
    top_indices = np.argsort(similarities)[-k:][::-1]

    # Return the top-k most relevant text chunks
    return [text_chunks[i] for i in top_indices]

def generate_response(system_prompt, user_message, model="meta-llama/Llama-3.2-3B-Instruct"):
    """
    Generates a response from the AI model based on the system prompt and user message.

    Args:
    system_prompt (str): The system prompt to guide the AI's behavior.
    user_message (str): The user's message or query.
    model (str): The model to be used for generating the response. Default is "meta-llama/Llama-2-7B-chat-hf".

    Returns:
    dict: The response from the AI model.
    """
    response = client.chat.completions.create(
        model=model,
        temperature=0,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message}
        ]
    )
    return response

# 完整调用流程
def semantic_chunking_rag_pipeline(pdf_path, query):
    # 1. 提取PDF文本
    extracted_text = extract_text_from_pdf(pdf_path)

    # 2. 按句子分割
    sentences = extracted_text.split(". ")

    # 3. 生成句子嵌入
    embeddings = [get_embedding(sentence) for sentence in sentences]

    # 4. 计算句子间相似度
    similarities = [cosine_similarity(embeddings[i], embeddings[i + 1]) for i in range(len(embeddings) - 1)]

    # 5. 计算断点（使用百分位数方法）
    breakpoints = compute_breakpoints(similarities, method="percentile", threshold=90)

    # 6. 分割成语义块
    text_chunks = split_into_chunks(sentences, breakpoints)

    # 7. 创建块嵌入
    chunk_embeddings = create_embeddings(text_chunks)

    # 8. 语义搜索
    top_chunks = semantic_search(query, text_chunks, chunk_embeddings, k=2)

    # 9. 生成回答
    system_prompt = "You are an AI assistant that strictly answers based on the given context. If the answer cannot be derived directly from the provided context, respond with: 'I do not have enough information to answer that.'"
    user_prompt = "\n".join([f"Context {i + 1}:\n{chunk}\n=====================================\n" for i, chunk in enumerate(top_chunks)])
    user_prompt = f"{user_prompt}\nQuestion: {query}"

    ai_response = generate_response(system_prompt, user_prompt)
    return ai_response.choices[0].message.content

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

class EncoderDecoder(nn.Module):
	"""
	标准的Encoder-Decoder架构。
	"""
	def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
		super(EncoderDecoder,self).__init__()
		self.encoder = encoder
		self.decoder = decoder
		# 源语言和目标语言的embedding
		self.src_embed = src_embed
		self.tgt_embed = tgt_embed
		# generator主要是根据Decoder的隐状态输出当前时刻的词(单个词)
		# 基本的实现就是隐状态输入一个全连接层，然后接一个softmax变成概率
		self.generator = generator

	def forward(self, src, tgt, src_mask, tgt_mask):
		# 首先调用encode方法对输入进行编码，然后调用decode方法进行解码
		return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)

	def encode(self, src, src_mask):
		# 调用self.encoder函数
		return self.encoder(self.src_embed(src), src_mask)

	def decode(self, memory, src_mask, tgt, tgt_mask):
		# 调用self.decoder函数 注意⚠️：这里定义的memery是encoder的输出结果。
		return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

1
2
3
4
5
6
7
8

class Generator(nn.Module):
	def __init__(self, d_model, vocab):
		super(Generator, self).__init__()
		# d_model是Decoder输出的大小，vocab是词典的大小
		self.proj = nn.Linear(d_model, vocab)

	def forward(self, x):
		return F.log_softmax(self.proj(x), dim=-1)

1
2
3
4
5
6

criterion = nn.CrossEntropyLoss()

x = torch.randn(1,5) # 服从0-1的正太分布。
y = torch.empty(1, dtype = torch.long).random_(5)

loss = criterion(x,y)

1
2
3
4
5

m = F.log_softmax(x, dim=1)
criterion = nn.NLLLoss()
x = torch.randn(1, 5)
y = torch.empty(1, dtype = torch.long).random_(5)
loss = criterion(m(x), y)

1
2
3

def clones(module, N):
	# 克隆N个完全相同的SubLayer，使用了copy.deepcopy
	return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

class Encoder(nn.Module):
	# Encoder是N个EncoderLayer的stack
	def __init__(self, layer, N):
		super(Encoder, self).__init__()
		# layer是一个SubLayer，我们clone N个
		self.layers = clones(layer, N)
		# 再加一个LayerNorm层
		self.norm = LayerNorm(layer.size)

	def forward(self, x, mask):
		# 逐层进行处理
		for layer in self.layers:
			x = layer(x, mask)
		# 最后进行LayerNorm
		return self.norm(x)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

class LayerNorm(nn.Module):
	def __init__(self, features, eps=1e-6):
		super(LayerNorm, self).__init__()
		self.a_2 = nn.Parameter(torch.ones(features))
		self.b_2. = nn.Parameter(torch.zeros(feagures))
		self.eps = eps
	def forward(self, x):
		mean = x.mean(-1, keepdim = True)
		std = x.std(-1, keepdim = True)
		return self.a_2 * (x - mean)/(std+self.eps) + self.b_2

1
2
3

x -> attention(x) -> x+self-attention(x)[残差] -> layernorm(x+self-attention(x)) => y

y -> dense(y) -> y+dense(y) -> layernorm(y+dense(y)) => z(输入下一层)

1
2

x -> layernorm(x) -> attention(layernorm(x)) -> a + attention(layernorm(x)) => y
y -> layernorm(y) -> dense(layernorm(y)) -> y+dense(layernorm(y))

1
2
3
4

def forward(self, x, mask):
	for layer in self.layers:
		x = layer(x, mask)
	return self.norm(x)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

class SublayerConnection(nn.Module):
	"""
	LayerNorm+sublayer(Self-Attention/Dense) + dropout + 残差连接
	为了简单，把LayerNorm放到了前面，这和原始论文稍有不同，原始论文LayerNorm在最后。
	"""
	def __init__(self, size, dropout):
		supper(SublayerConnection, self).__init__()
		self.norm = LayerNorm(size)
		self.dropout = nn.Droupout(dropout)
	def forward(self, x, sublayer):
		# sublayer是传入的参数,之后进行残差连接
		return x+self.dropout(sublayer(self.norm(x)))

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

class EncoderLayer(nn.Module):
	# EncoderLayer由self-attn和feed_forward组成
	def __init__(self, size, self_attn, feed_forward, dropout):
		super(EncoderLayer, self).__init__()
		self.size = size
		self.self_attn = self_attn
		self.feed_forward = feed_forward
		self.sublayer = clones(SublayerConnection(size, dropout), 2)
	def forward(self, x, mask):
		x = self.sublayer[0](x, lambda x:self.self_attn(x,x,x,mask))
		return self.sublayer[1](x, self.feed_forward)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

class Decoder(nn.Module):
	def __init__(self, layer, N):
		super(Decoder, self).__init__()
		self.layers = clones(layer, N)
		self.norm = LayerNorm(layer.size)

	def forward(self, x, memory, src_mask, tgt_mask):
		for layer in self.layers:
			x = layer(x, memory, src_mask, tgt_mask)
		return self.norm(x)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

class DecoderLayer(nn.Module):
	# Decoder包括self-attn, src-attn, feed_forward
	def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
		super(DecoderLayer, self).__init__()
		self.size = size
		self.self_attn = self_attn
		self.src_attn = src_attn
		self.feed_forward = feed_forward
		self.sublayer = clones(SublayerConnection(size, dropout), 3)
	def forward(self, x, memory, src_mask, tgt_mask):
		m = memory
		x = self.sublayer[0](x, lambda x: self.self_attn(x,x,x,tgt_mask))
		x = self.sublayer[1](x, lambda x: self.src_attn(x,m,m,src_mask))
		return self.sublayer[2](x, self.feed_forward)

1
2
3
4
5

def subsequent_mask(size):
	# mask out subsequent positoins
	attn_shape = (1, size, size)
	subsequent_mask = np.triu(np.ones(attn_shape),k=1).astype('uint8')
	return torch.from_numpy(subsequent_mask) == 0

1
2
3
4
5
6
7

print(subsequent_mask(5))
# 输出
1 0 0 0 0
1 1 0 0 0
1 1 1 0 0
1 1 1 1 0
1 1 1 1 1

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

def attention(query, key, value, mask=None,dropout=None):
	d_k = query.size(-1)
	scores = torch.matmul(query,key.transpose(-2,-1))/math.sqrt(d_k)
	if mask is not None:
		scores = scores.mask_fill(mask==0,-1e9)
	p_attn = F.softmax(scores, dim=-1)

	if dropout is not None:
		p_attn = dropout(p_attn)
	return torch.matmul(p_attn,value),p_attn

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

class MultiHeadAttention(nn.Module):
	def __init__(self, h, d_model, dropout=0.1):
		super(MultiHeadAttention, self).__init__()
		assert d_model % h == 0
		self.d_k = d_model // h # 这里是整除
		self.h = h
		self.linears = clones(nn.Linear(d_model,d_k),4)
		self.attn = None
		self.dropout = nn.Dropout(p=dropout)

	def foward(self, query, key, value, mask=None):
		if mask is not None:
			# 所有h个head的mask都是相同的
			mask = mask.unsqueeze(1)
		nbatches = query.size(0)
		# 1) 首先使用线性变换，然后把d_model分配给h个Head，每个head为d_k=d_model/h
		query, key, value = [l(x).view(nbatches,-1,self.h,self.d_k).transpose(1,2) for l,x in zip(self.linears,(query, key, value))]
		# 2) 使用attention函数计算
		x, self.attn = attention(query, key, value, mask=mask,dropout = self.dropout)
		# 3）
		x = x.transpose(1,2).contiguous().view(nbatches, -1,self.h*self.d_k)
		return self.linears[-1](x)