Spaces:
Sleeping
Sleeping
| import faiss | |
| import pickle | |
| import threading | |
| import time | |
| import torch | |
| import numpy as np | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| from rank_bm25 import BM25Okapi | |
| from sentence_transformers import SentenceTransformer | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| class FinancialChatbot: | |
| def __init__(self): | |
| # Load FAISS index | |
| self.faiss_index = faiss.read_index("financial_faiss.index") | |
| with open("index_map.pkl", "rb") as f: | |
| self.index_map = pickle.load(f) | |
| # Load BM25 keyword-based search | |
| with open("bm25_corpus.pkl", "rb") as f: | |
| self.bm25_corpus = pickle.load(f) | |
| self.bm25 = BM25Okapi(self.bm25_corpus) | |
| # Load SentenceTransformer for embedding-based retrieval | |
| self.sbert_model = SentenceTransformer("all-MiniLM-L6-v2") | |
| # Load Qwen Model | |
| model_name = "Qwen/Qwen2.5-1.5b" | |
| self.qwen_model = AutoModelForCausalLM.from_pretrained( | |
| model_name, torch_dtype="auto", device_map="auto", trust_remote_code=True | |
| ) | |
| self.qwen_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) | |
| # Guardrail: Blocked Words | |
| self.BLOCKED_WORDS = [ | |
| "hack", "bypass", "illegal", "exploit", "scam", "kill", "laundering", | |
| "murder", "suicide", "self-harm", "assault", "bomb", "terrorism", | |
| "attack", "genocide", "mass shooting", "credit card number" | |
| ] | |
| # Relevance threshold | |
| self.min_similarity_threshold = 0.2 | |
| def moderate_query(self, query): | |
| """Check if the query contains inappropriate words.""" | |
| query_lower = query.lower() | |
| for word in self.BLOCKED_WORDS: | |
| if word in query_lower: | |
| return False # Block query | |
| return True # Allow query | |
| def query_faiss(self, query, top_k=5): | |
| """Retrieve relevant documents using FAISS and compute confidence scores.""" | |
| query_embedding = self.sbert_model.encode([query], convert_to_numpy=True) | |
| distances, indices = self.faiss_index.search(query_embedding, top_k) | |
| results = [] | |
| confidence_scores = [] | |
| for idx, dist in zip(indices[0], distances[0]): | |
| if idx in self.index_map: | |
| similarity = 1 / (1 + dist) # Convert L2 distance to similarity | |
| results.append(self.index_map[idx]) | |
| confidence_scores.append(similarity) | |
| return results, confidence_scores | |
| def query_bm25(self, query, top_k=5): | |
| """Retrieve relevant documents using BM25 keyword-based search.""" | |
| tokenized_query = query.lower().split() | |
| scores = self.bm25.get_scores(tokenized_query) | |
| top_indices = np.argsort(scores)[::-1][:top_k] | |
| results = [] | |
| confidence_scores = [] | |
| for idx in top_indices: | |
| if scores[idx] > 0: # Ignore zero-score matches | |
| results.append(self.bm25_corpus[idx]) | |
| confidence_scores.append(scores[idx]) | |
| return results, confidence_scores | |
| def generate_answer(self, context, question): | |
| """Generate answer using the Qwen model.""" | |
| input_text = f"Context: {context}\nQuestion: {question}\nAnswer:" | |
| inputs = self.qwen_tokenizer.encode(input_text, return_tensors="pt") | |
| outputs = self.qwen_model.generate(inputs, max_length=100) | |
| return self.qwen_tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| def get_answer(self, query, timeout=200): | |
| """Fetch an answer from FAISS and Qwen model while handling timeouts.""" | |
| result = ["No relevant information found", 0.0] # Default response | |
| def task(): | |
| if query.lower() in ["hi", "hello", "hey"]: | |
| result[:] = ["Hi, how can I help you?", 1.0] | |
| return | |
| if not self.moderate_query(query): | |
| result[:] = ["I'm unable to process your request due to inappropriate language.", 1.0] | |
| return | |
| faiss_results, faiss_conf = self.query_faiss(query) | |
| bm25_results, bm25_conf = self.query_bm25(query) | |
| all_results = faiss_results + bm25_results | |
| all_conf = faiss_conf + bm25_conf | |
| # Check relevance | |
| if not all_results or max(all_conf, default=0) < self.min_similarity_threshold: | |
| result[:] = ["No relevant information found", 1.0] | |
| return | |
| context = " ".join(all_results) | |
| answer = self.generate_answer(context, query) | |
| last_index = answer.rfind("Answer") | |
| if answer[last_index+9:11] == "--": | |
| result[:] = ["No relevant information found", 1.0] | |
| else: | |
| result[:] = [answer[last_index:], max(all_conf, default=0.9)] | |
| thread = threading.Thread(target=task) | |
| thread.start() | |
| thread.join(timeout) | |
| if thread.is_alive(): | |
| return "No relevant information found", 1.0 # Timeout case | |
| return tuple(result) | |