import os import torch import pandas as pd from transformers import AutoTokenizer, AutoModel, pipeline from sentence_transformers import SentenceTransformer, CrossEncoder from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain_community.document_loaders import PyPDFLoader, CSVLoader from langchain_community.vectorstores import FAISS from langchain.prompts import PromptTemplate from pathlib import Path from langchain_huggingface import HuggingFaceEmbeddings from langchain_community.document_loaders import DirectoryLoader from datetime import datetime from concurrent.futures import ThreadPoolExecutor class RagWithScore: def __init__(self, model_name="sentence-transformers/all-MiniLM-L6-v2", cross_encoder_name="cross-encoder/ms-marco-TinyBERT-L-2-v2", llm_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0", documents_dir="financial_docs"): """ Initialize the Financial RAG system Args: model_name: The embedding model name cross_encoder_name: The cross-encoder model for reranking llm_name: Small language model for generation documents_dir: Directory containing financial statements """ # Initialize embedding model self.embedding_model = HuggingFaceEmbeddings(model_name=model_name) # Initialize cross-encoder for reranking self.cross_encoder = CrossEncoder(cross_encoder_name) # Initialize small language model self.tokenizer = AutoTokenizer.from_pretrained(llm_name) self.llm = pipeline( "text-generation", model=llm_name, tokenizer=self.tokenizer, torch_dtype=torch.bfloat16, device_map="auto", max_new_tokens=512, do_sample=False, # Set to False for deterministic outputs temperature=0.2, # Reduce randomness top_p=1.0 # No nucleus sampling ) # Store paths self.documents_dir = documents_dir self.vector_store = None # Input guardrail rules - sensitive terms/patterns self.guardrail_patterns = [ "insider trading", "stock manipulation", "fraud detection", "embezzlement", "money laundering", "tax evasion", "illegal activities" ] # Confidence score thresholds self.confidence_thresholds = { "high": 0.75, "medium": 0.5, "low": 0.3 } import os ## Loadung document and creating vector index at the start of the application def load_and_process_documents(self): """Load, split and process financial documents""" print("Processing documents to create FAISS index...") loader = DirectoryLoader('./financial_docs', glob="**/*.pdf") documents = loader.load() # Split documents into chunks text_splitter = RecursiveCharacterTextSplitter( chunk_size=1000, chunk_overlap=200 ) chunks = text_splitter.split_documents(documents) print(len(chunks)) # Create and save FAISS vector store self.vector_store = FAISS.from_documents(chunks, embedding=self.embedding_model) self.vector_store.save_local("faiss_index") return self.vector_store ## generating response with the query and context by the help of the prompt and calling the slm with the prompt def generate_answer(self, query, context): """Generate answer and calculate confidence score concurrently.""" # Format context into a single string context_str = "\n\n".join([doc.page_content for doc in context]) # Define the prompt prompt = f""" You are a financial analyst assistant that helps answer questions about company financial statements. Use the provided financial information to give accurate and helpful answers. Context: {context_str} Question: {query} Instructions: 1. Be concise and avoid repetition. 2. If the question is too broad or unclear, ask for clarification or provide a general overview. 3. Only if the question requires numerical calculations, show the step-by-step logic behind the calculation before providing the final answer. 4. If the context is insufficient, say "Not enough information to answer this question." 5. Do not provide sources unless explicitly asked. Answer: """ # Generate answer using the language model response = self.llm(prompt)[0]['generated_text'] answer = response[len(prompt):].strip() # # Calculate confidence score concurrently # with ThreadPoolExecutor() as executor: # future_confidence = executor.submit(self.calculate_confidence_score, query, context, answer) # confidence_score = future_confidence.result() return answer ## for confidence score cosine similarity is calculated between the query embedding and answer embedding def calculate_confidence_score(self, query, retrieved_docs, answer): """ Calculate confidence score using embedding similarity (parallelized). """ # Get embeddings for query and answer query_embedding = self.embedding_model.embed_query(query) answer_embedding = self.embedding_model.embed_query(answer) # Calculate cosine similarity dot_product = sum(a * b for a, b in zip(query_embedding, answer_embedding)) magnitude_a = sum(a * a for a in query_embedding) ** 0.5 magnitude_b = sum(b * b for b in answer_embedding) ** 0.5 similarity = dot_product / (magnitude_a * magnitude_b) if magnitude_a * magnitude_b > 0 else 0 return similarity ## confidence level is determined from the confidence score def get_confidence_level(self, confidence_score): """ Convert numerical confidence score to a level (high, medium, low) Args: confidence_score: Float between 0 and 1 Returns: str: Confidence level ("high", "medium", or "low") """ if confidence_score >= self.confidence_thresholds["high"]: return "high" elif confidence_score >= self.confidence_thresholds["medium"]: return "medium" elif confidence_score >= self.confidence_thresholds["low"]: return "low" else: return "very low" ## guardrail is applied to filter harmful user queries def apply_input_guardrail(self, query): """Check if query violates input guardrails""" query_lower = query.lower() for pattern in self.guardrail_patterns: if pattern in query_lower: return True, f"I cannot process queries about {pattern}. Please reformulate your question." return False, "" ## first the to 5 chunks are retrieved. then after reranking with cross encoder top 2 are rerieved def retrieve_with_reranking(self, query, top_k=5, rerank_top_k=3): print("retrieve_with_reranking start") print(datetime.now().strftime('%Y-%m-%d %H:%M:%S')) """Retrieve relevant chunks and rerank them with cross-encoder""" # Initial retrieval using embedding similarity docs_and_scores = self.vector_store.similarity_search_with_score(query, k=top_k) # Sort retrieved documents by FAISS similarity score (deterministic sorting) docs_and_scores.sort(key=lambda x: x[1], reverse=True) # Prepare pairs for cross-encoder pairs = [(query, doc.page_content) for doc, _ in docs_and_scores] print(pairs) print(len(pairs)) # Get scores from cross-encoder scores = self.cross_encoder.predict(pairs) print(scores) # Sort by cross-encoder scores reranked_results = sorted(zip(docs_and_scores, scores), key=lambda x: x[1], reverse=True) print(reranked_results) print(len(reranked_results)) # Return the top reranked results print("retrieve_with_reranking end") print(datetime.now().strftime('%Y-%m-%d %H:%M:%S')) return [doc for (doc, _), _ in reranked_results[:rerank_top_k]] ## to handle irrerelevant questions, a rule based claasifier is bein used to classify the questions def is_financial_question(self,query): financial_keywords = [ "finance", "financial", "revenue", "profit", "loss", "ebitda", "cash flow", "balance sheet", "income statement", "stock", "bond", "investment", "risk", "interest rate", "inflation", "debt", "equity", "valuation", "dividend", "market", "economy", "GDP", "currency", "exchange rate", "tax", "audit", "compliance", "regulation", "SEC", "earnings", "capital", "asset", "liability" ] query_lower = query.lower() return any(keyword in query_lower for keyword in financial_keywords) ##the pipeline of answer and confidence score generation from the query def answer_question(self, query): """End-to-end pipeline to answer a question with confidence score""" # Apply input guardrail blocked, message = self.apply_input_guardrail(query) if blocked: return {"answer": message, "source_documents": [], "blocked": True, "confidence_score": 0, "confidence_level": "none"} if not self.is_financial_question(query): return { "answer": "This question is outside the scope of financial data. Please ask a question related to finance.", "source_documents": [], "blocked": True, "confidence_score": 0, "confidence_level": "none" } # Retrieve and rerank relevant contexts reranked_docs = self.retrieve_with_reranking(query) # Generate answer and confidence score answer = self.generate_answer(query, reranked_docs) print(f"answer : {answer}") # Calculate confidence score confidence_score = self.calculate_confidence_score(query, reranked_docs, answer) print(f"confidence score : {confidence_score}") confidence_level = self.get_confidence_level(confidence_score) return { "answer": answer, "source_documents": reranked_docs, "blocked": False, "confidence_score": confidence_score, "confidence_level": confidence_level }