Spaces:
Running
Running
File size: 2,418 Bytes
509d6f6 bb4ea25 509d6f6 bb4ea25 509d6f6 bb4ea25 509d6f6 bb4ea25 509d6f6 bb4ea25 509d6f6 bb4ea25 509d6f6 | 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 | # -*- 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) |