Initial commit

2025-09-27 12:02:22 +08:00
commit a4443765ee
33 changed files with 3896 additions and 0 deletions
@@ -0,0 +1,88 @@
+import torch
+import sqlite3
+import numpy as np
+from torch import Tensor
+from typing import Dict, List, Tuple
+
+
+class Retriever:
+    def __init__(self, db_path=None):
+        self.data: Dict[str, Tensor] = {}
+        self.embedding_cache: Tensor = None
+        self.is_caculated: bool = False
+        
+        if db_path is not None:
+            self.load(db_path)
+                
+    def retrieve(self, query: Tensor, top_k: int) -> List[Tuple[str, float]]:
+        if not self.data:
+            return []
+        
+        query = query.flatten().unsqueeze(1)                            # [dim, 1]
+        norm_embeddings = self._embeddings.to(
+            device=query.device,
+            dtype=query.dtype
+        )                                                               # [n_vectors, dim]
+        sim_scores = torch.matmul(norm_embeddings, query).squeeze()     # [n_vectors]
+        
+        top_k = min(top_k, len(self.data))
+        indices = sim_scores.topk(top_k).indices
+        keys = list(self.data.keys())
+        
+        return [(keys[i], sim_scores[i].item()) for i in indices]
+    
+    def add_vector(self, key: str, vector_data: Tensor):
+        self.is_caculated = False
+        self.data[key] = vector_data.flatten().float().cpu()
+        
+    def delete_vector(self, key: str):
+        self.is_caculated = False
+        self.data.pop(key, None)
+    
+    def save(self, db_path):
+        conn = sqlite3.connect(db_path)
+        cursor = conn.cursor()
+        self._init_db(cursor)
+        cursor.execute('DELETE FROM vectors')
+        
+        for item, vec in self.data.items():
+            vec_bytes = vec.numpy().tobytes()
+            cursor.execute('INSERT OR REPLACE INTO vectors (key, vector) VALUES (?, ?)', 
+                           (item, vec_bytes))
+        
+        conn.commit()
+        conn.close()
+
+    def load(self, db_path):
+        conn = sqlite3.connect(db_path)
+        cursor = conn.cursor()
+        self._init_db(cursor)
+        cursor.execute('SELECT key, vector FROM vectors')
+        rows = cursor.fetchall()
+        self.data = {}
+        
+        for row in rows:
+            key, vec_bytes = row
+            vec_numpy = np.frombuffer(vec_bytes, dtype=np.float32).copy()
+            vec = torch.from_numpy(vec_numpy)
+            self.data[key] = vec
+        
+        conn.close()
+        
+    def _init_db(self,cursor: sqlite3.Cursor):
+        # Create table if not exists (in case loading from a new database)
+        cursor.execute('''
+            CREATE TABLE IF NOT EXISTS vectors (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                key TEXT UNIQUE NOT NULL,
+                vector BLOB NOT NULL
+            )''')
+    
+    @property
+    def _embeddings(self) -> Tensor:
+        if not self.is_caculated:
+            embeddings = torch.stack(list(self.data.values())) 
+            norm_embeddings = embeddings / torch.norm(embeddings, dim=-1, keepdim=True)
+            self.embedding_cache = norm_embeddings
+        
+        return self.embedding_cache
@@ -0,0 +1,127 @@
+import re
+import torch
+import torch.nn.functional as F
+
+from abc import ABC, abstractmethod
+from torch import Tensor
+from typing import List, Callable, Optional
+
+
+class BaseTextSplitter(ABC):
+    def __init__(
+        self,
+        max_len: int = 512,
+        chunk_overlap: int = 0,
+    ):
+        if max_len <= 0:
+            raise ValueError("max_len must be > 0")
+        if chunk_overlap < 0:
+            raise ValueError("chunk_overlap must be >= 0")
+
+        self.max_len = max_len
+        self.chunk_overlap = chunk_overlap
+
+    @abstractmethod
+    def split(self, text: str, **kwargs) -> List[str]:
+        raise NotImplementedError
+
+    def preprocess(self, text: str) -> str:
+        return text.strip()
+
+    def postprocess(self, chunks: List[str]) -> List[str]:
+        return [chunk.strip() for chunk in chunks if chunk.strip()]
+
+
+class PriorityTextSplitter(BaseTextSplitter):
+    def __init__(
+        self,
+        separators: List[str],
+        max_len: int = 512,
+        chunk_overlap: int = 0,
+    ):
+        super().__init__(max_len=max_len, chunk_overlap=chunk_overlap)
+        if not separators:
+            raise ValueError("separators must be a non-empty list")
+        self.separators = separators
+
+    def split(self, text: str) -> List[str]:
+        text = self.preprocess(text)
+        for sep in self.separators:
+            parts = text.split(sep)
+
+            valid_parts = [p.strip() for p in parts if p.strip()]
+            if len(valid_parts) > 1:
+                return self.postprocess(valid_parts)
+        return [text]
+
+
+class SemanticTextSplitter(BaseTextSplitter):
+
+    DEFAULT_PATTERN = r'(?<=[。！？!?])(?=(?:[^"\'‘’“”]*["\'‘’“”][^"\'‘’“”]*["\'‘’“”])*[^"\'‘’“”]*$)'
+
+    def __init__(
+        self,
+        embedding_func: Callable[[List[str]], List[Tensor]],
+        pattern: Optional[str] = None,
+        max_len: int = 512,
+        chunk_overlap: int = 0,
+    ):
+        super().__init__(max_len=max_len, chunk_overlap=chunk_overlap)
+        if not callable(embedding_func):
+            raise TypeError("embedding_func must be callable")
+        self.embedding_func = embedding_func
+        self.pattern = pattern or SemanticTextSplitter.DEFAULT_PATTERN
+
+    def split(
+        self,
+        text: str,
+        threshold: float = 0.5,
+        window_size: int = 1,
+    ) -> List[str]:
+        text = self.preprocess(text)
+        sentences = [s.strip() for s in re.split(self.pattern, text) if s.strip()]
+
+        if len(sentences) <= 1:
+            return self.postprocess(sentences)
+
+        try:
+            sentence_embs = self.embedding_func(sentences)
+        except Exception as e:
+            raise RuntimeError(f"Embedding generation failed: {e}")
+
+        if len(sentence_embs) != len(sentences):
+            raise ValueError("Embedding function must return one vector per sentence")
+
+        chunks = []
+        emb_tensor = torch.stack(sentence_embs)  # shape: [N, D]
+        current_chunk: List[str] = [sentences[0]]
+
+        for i in range(1, len(sentences)):
+            start_prev = max(0, i - window_size)
+            end_prev = i
+            start_next = i
+            end_next = min(len(sentences), i + window_size)
+            
+            prev_window_emb = emb_tensor[start_prev:end_prev].mean(dim=0)
+            next_window_emb = emb_tensor[start_next:end_next].mean(dim=0)
+
+            similarity = F.cosine_similarity(
+                prev_window_emb.unsqueeze(0),
+                next_window_emb.unsqueeze(0),
+                dim=1
+            ).item()
+
+            dynamic_threshold = max(threshold * (1 - 0.03 * (end_next - start_prev)), 0.2)
+
+            if similarity < dynamic_threshold:
+                chunks.append(" ".join(current_chunk))
+                overlap_start = max(0, len(current_chunk) - self.chunk_overlap)
+                current_chunk = current_chunk[overlap_start:]
+                current_chunk.append(sentences[i])
+            else:
+                current_chunk.append(sentences[i])
+
+        if current_chunk:
+            chunks.append(" ".join(current_chunk))
+
+        return self.postprocess(chunks)