Spaces:
Running
Running
| # -*- coding: utf-8 -*- | |
| from typing import List | |
| import numpy as np | |
| from config import APP_CONFIG | |
| try: | |
| from sklearn.feature_extraction.text import TfidfVectorizer | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| SKLEARN_AVAILABLE = True | |
| except ImportError: | |
| SKLEARN_AVAILABLE = False | |
| class ResearchRAGEngine: | |
| def __init__(self, raw_texts: List[str]): | |
| self.chunks = [] | |
| self.vectorizer = None | |
| self.tfidf_matrix = None | |
| self.chunk_size = APP_CONFIG["rag_chunk_size"] | |
| self._initialize_chunks(raw_texts) | |
| def _initialize_chunks(self, raw_texts: List[str]): | |
| for text in raw_texts: | |
| for i in range(0, len(text), self.chunk_size): | |
| chunk = text[i:i + self.chunk_size].strip() | |
| if len(chunk) > 50: | |
| self.chunks.append(chunk) | |
| if SKLEARN_AVAILABLE and self.chunks: | |
| self.vectorizer = TfidfVectorizer(stop_words='english', max_df=0.85) | |
| try: | |
| self.tfidf_matrix = self.vectorizer.fit_transform(self.chunks) | |
| except ValueError: | |
| self.vectorizer = None | |
| def get_relevant_context(self, query: str, top_k: int = None) -> str: | |
| if not self.chunks: | |
| return "" | |
| if top_k is None: | |
| top_k = APP_CONFIG["rag_top_k"] | |
| relevant_chunks = [] | |
| if self.vectorizer and self.tfidf_matrix is not None: | |
| query_vec = self.vectorizer.transform([query]) | |
| similarities = cosine_similarity(query_vec, self.tfidf_matrix).flatten() | |
| top_indices = similarities.argsort()[-top_k:][::-1] | |
| relevant_chunks = [self.chunks[i] for i in top_indices if similarities[i] > 0.05] | |
| else: | |
| scores = [] | |
| query_words = set(query.lower().split()) | |
| for chunk in self.chunks: | |
| chunk_lower = chunk.lower() | |
| score = sum(1 for word in query_words if word in chunk_lower) | |
| scores.append(score) | |
| top_indices = np.argsort(scores)[-top_k:] | |
| relevant_chunks = [self.chunks[i] for i in top_indices if scores[i] > 0] | |
| if not relevant_chunks: | |
| return "" | |
| return "\n--- Target Context Snippet ---\n" + "\n\n".join(relevant_chunks) | |
| def initialize_rag(texts: List[str]): | |
| return ResearchRAGEngine(texts) |