Spaces:

ShadowTEM
/

Chatbot

Sleeping

Chatbot / app /vectorstore.py

Upgrading chunks from 700 to 1200

89fe4c9 verified over 1 year ago

13.8 kB

# vectorstore.py
import os
import json
from datetime import datetime
# import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from pypdf import PdfReader
import requests
from bs4 import BeautifulSoup

# Updated LangChain imports
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings


class vectorstore:
    # 1 is pdf, 2 is website, and 3 is an already created vectorstore
    def __init__(self, path, Initlize_with=3):
        if Initlize_with == 1:
            self.index_path = self.create_vectorstore_from_pdf(path, embedder_model="all-MiniLM-L6-v2")
            self.vectorstore = self.load_vectorstore(self.index_path)
        elif Initlize_with == 2:
            self.index_path = self.create_vectorstore_from_website(path, embedder_model="all-MiniLM-L6-v2")
            self.vectorstore = self.load_vectorstore(self.index_path)
        elif Initlize_with == 3:
            self.index_path = path
            self.vectorstore = self.load_vectorstore(self.index_path)
            
    def chunk_text(self,text, chunk_size=1200):
        """
        Split text into chunks of approximately 'chunk_size' words.
        
        Parameters:
            text (str): The input text.
            chunk_size (int): Maximum number of words per chunk.
        
        Returns:
            List[str]: A list of text chunks.
        """
        words = text.split()
        chunks = []
        for i in range(0, len(words), chunk_size):
            chunk = " ".join(words[i:i+chunk_size])
            chunks.append(chunk)
        return chunks

    def create_vectorstore_from_pdf(self,pdf_path, chunk_size=1200, embedder_model="all-MiniLM-L6-v2"):
        """
        Extract text from a PDF using pypdf, chunk the text, compute embeddings,
        and create a FAISS index.
        
        Parameters:
            pdf_path (str): Path to the PDF file.
            chunk_size (int): Number of words per chunk.
            embedder_model (str): The sentence-transformer model to use.
            
        Returns:
            index: A FAISS index containing the embeddings.
            chunks: A list of text chunks.
            embedder: The SentenceTransformer embedder.
        """
        # Load PDF using LangChain's PyPDFLoader
        loader = PyPDFLoader(pdf_path)
        documents = loader.load()
        text = " ".join([doc.page_content for doc in documents])
        
        # Chunk text using existing method (word-based)
        chunks = self.chunk_text(text, chunk_size)
        
        # Create embeddings and FAISS vectorstore with LangChain
        embedder = HuggingFaceEmbeddings(model_name=embedder_model)
        vectorstore = FAISS.from_texts(chunks, embedder)
        
        # Save the vectorstore
        index_path = self.save_vectorstore_with_timestamp_and_without(vectorstore, embedder_model)
        return index_path

    def create_vectorstore_from_website(self, url, chunk_size=1200, embedder_model="all-MiniLM-L6-v2"):
        """
        Fetch text from a website, chunk the text, compute embeddings,
        and create a FAISS index.
        
        Parameters:
            url (str): The URL of the website.
            chunk_size (int): Number of words per chunk.
            embedder_model (str): The sentence-transformer model to use.
            
        Returns:
            index: A FAISS index containing the embeddings.
            chunks: A list of text chunks.
            embedder: The SentenceTransformer embedder.
        """
        # Load website using LangChain's WebBaseLoader
        loader = WebBaseLoader(url)
        documents = loader.load()
        text = " ".join([doc.page_content for doc in documents])
        
        # Chunk text using existing method
        chunks = self.chunk_text(text, chunk_size)
        
        # Create embeddings and FAISS vectorstore
        embedder = HuggingFaceEmbeddings(model_name=embedder_model)
        vectorstore = FAISS.from_texts(chunks, embedder)
        
        # Save the vectorstore
        index_path = self.save_vectorstore_with_timestamp_and_without(vectorstore, embedder_model)
        return index_path

    def add_to_vectorstore_web(self, url, index_file_path, chunk_size=1200):
        """
        Fetch text from a website, chunk the text, compute embeddings using the existing embedder,
        and add them to the existing FAISS vectorstore.
        Parameters:
            url (str): The URL of the website.
            index_file_path (str): The file path of the saved vectorstore index.
            chunk_size (int): Maximum number of words per chunk.
        Returns:
            str: The updated vectorstore index file path.
        """
    # Load the existing vectorstore
        vectorstore = self.load_vectorstore(index_file_path)
        
        # Fetch and extract text from the website
        loader = WebBaseLoader(url)
        documents = loader.load()
        text = " ".join([doc.page_content for doc in documents])
        
        # Chunk the text
        new_chunks = self.chunk_text(text, chunk_size)
        
        # Add new chunks to the vectorstore
        vectorstore.add_texts(new_chunks)
        
        # Save the updated vectorstore
        updated_index_path = self.save_vectorstore_with_timestamp_and_without(vectorstore, vectorstore.embedding_function.model_name)
        return updated_index_path


    def add_to_vectorstore_from_pdf(self, pdf_path, chunk_size=1200):
        """
        Extract text from a PDF using PyPDF2, chunk the text, compute embeddings using the existing embedder,
        and add them to the existing FAISS vectorstore.
        Parameters:
            pdf_path (str): Path to the PDF file.
            index_file_path (str): The file path of the saved vectorstore index.
            chunk_size (int): Maximum number of words per chunk.
        Returns:
            str: The updated vectorstore index file path.
        """
    # Extract text from PDF
        loader = PyPDFLoader(pdf_path)
        documents = loader.load()
        text = " ".join([doc.page_content for doc in documents])
        
        # Chunk the text
        new_chunks = self.chunk_text(text, chunk_size)
        
        # Add new chunks to the existing vectorstore
        self.vectorstore.add_texts(new_chunks)
        
        # Save the updated vectorstore
        self.index_path = self.save_vectorstore_with_timestamp_and_without(self.vectorstore, self.vectorstore.embedding_function.model_name)
        return self.index_path

    def save_vectorstore_with_timestamp_and_without(self, vectorstore, embedder_model=None):
        vector_db_folder = os.path.join(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")), "VectorDB")
        os.makedirs(vector_db_folder, exist_ok=True)
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        folder_name = f"vectorstore_{timestamp}"
        folder_path = os.path.join(vector_db_folder, folder_name)
        
        # Save the LangChain FAISS vectorstore
        vectorstore.save_local(folder_path)
        
        # Save metadata with embedder model name
        metadata = {"embedder_model": embedder_model or "all-MiniLM-L6-v2"}
        with open(os.path.join(folder_path, "metadata.json"), "w") as f:
            json.dump(metadata, f)
        
        # Save to main folder
        main_folder = os.path.join(vector_db_folder, "vectorstore_mainV2")
        vectorstore.save_local(main_folder)
        with open(os.path.join(main_folder, "metadata.json"), "w") as f:
            json.dump(metadata, f)
        
        return main_folder

    def load_vectorstore(self, index_path):
        # Load metadata to get embedder model
        metadata_path = os.path.join(index_path, "metadata.json")
        with open(metadata_path, "r") as f:
            metadata = json.load(f)
        embedder_model = metadata["embedder_model"]
        
        # Create embedder and load vectorstore
        embedder = HuggingFaceEmbeddings(model_name=embedder_model)
        vectorstore = FAISS.load_local(index_path, embedder,allow_dangerous_deserialization=True)
        return vectorstore
    
    
    def search_vectorstore(self, query, top_k=5):
        # Search using LangChain's similarity_search
        docs = self.vectorstore.similarity_search(query, k=top_k)
        results = [doc.page_content for doc in docs]
        return results

    def log_conversation(self, user_text, bot_text=""):
        # Existing JSON logging
        logs_folder = os.path.join(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")), "logs")
        os.makedirs(logs_folder, exist_ok=True)
        log_file = os.path.join(logs_folder, "conversation_logs.json")
        entry = {
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "User": {"text": user_text},
            "bot": {"text": bot_text}
        }
        if os.path.exists(log_file):
            with open(log_file, "r", encoding="utf-8") as f:
                try:
                    logs = json.load(f)
                except json.JSONDecodeError:
                    logs = []
        else:
            logs = []
        logs.append(entry)
        with open(log_file, "w", encoding="utf-8") as f:
            json.dump(logs, f, indent=4)
        
        # Add to LangMem memory
        self.memory.add_message({"role": "user", "content": user_text})
        if bot_text:
            self.memory.add_message({"role": "assistant", "content": bot_text})

# # ...existing code...

# import re
# import requests
# import nltk
# from nltk.tokenize import sent_tokenize
# import numpy as np

# # Download NLTK data (only once; consider moving this to a setup step)
# nltk.download('punkt')
# nltk.download('punkt_tab')

# class CancerDataIngestor:
#     def __init__(self,vector_db=None, source_urls=None, embedding_model=None):
#         """
#         Initialize with:
#         - source_urls: a dictionary mapping source names to URLs.
#         - embedding_model: a callable that takes text and returns an embedding.
#         - vector_db: an instance of your vector database.
#         """
#         self.sources = source_urls if source_urls is not None else {
#             'PMC': 'https://www.ncbi.nlm.nih.gov/pmc/',
#             'NCI': 'https://www.cancer.gov/about-cancer/understanding/statistics',
#             'WHO': 'https://www.who.int/cancer/en/',
#             'ClinicalTrials': 'https://clinicaltrials.gov/',
#             'Kaggle': 'https://www.kaggle.com/datasets'
#         }
#         # Use the provided embedding model or default to self.get_embedding
#         self.embedding_model = embedding_model if embedding_model is not None else self.get_embedding
#         if vector_db is None:
#             raise ValueError("A valid vector DB instance must be provided.")
#         self.vector_db = vector_db
#         # Initialize the SentenceTransformer model once.
#         self.model = SentenceTransformer("all-MiniLM-L6-v2")
    
#     def fetch_html(self, url):
#         """Fetch HTML content from a given URL."""
#         try:
#             response = requests.get(url)
#             response.raise_for_status()
#             return response.text
#         except Exception as e:
#             print(f"Error fetching URL {url}: {e}")
#             return None

#     def clean_html(self, html):
#         """Extract text from HTML and remove tags and extra spaces."""
#         if html is None:
#             return ""
#         from bs4 import BeautifulSoup  # in case not already imported above
#         soup = BeautifulSoup(html, 'html.parser')
#         for tag in soup(['script', 'style']):
#             tag.decompose()
#         text = soup.get_text(separator=' ')
#         text = re.sub(r'[^a-zA-Z0-9.,;:?!\s]', ' ', text)
#         text = re.sub(r'\s+', ' ', text)
#         return text.strip()

#     def preprocess_text(self, text):
#         """Lowercase the text and tokenize into sentences."""
#         text = text.lower()
#         sentences = sent_tokenize(text)
#         return sentences

#     def get_embedding(self, text):
#         """
#         Compute the embedding for a given text using the "all-MiniLM-L6-v2" model.
#         """
#         return self.model.encode(text)


  
#     def add_data_to_vectordb(self):
#         """
#         Iterate through all defined sources, fetch and clean the text,
#         generate embeddings for each sufficiently long sentence,
#         and add them to the vector DB using its underlying FAISS add_texts method.
#         """
#         for source, url in self.sources.items():
#             print(f"Processing source: {source}")
#             html_content = self.fetch_html(url)
#             cleaned_text = self.clean_html(html_content)
#             if not cleaned_text:
#                 print(f"No text fetched from {url}")
#                 continue
#             sentences = self.preprocess_text(cleaned_text)
#             texts = []
#             metadatas = []
#             for i, sentence in enumerate(sentences):
#                 if len(sentence) < 50:
#                     continue
#                 texts.append(sentence)
#                 metadatas.append({
#                     'source': source,
#                     'url': url,
#                     'sentence_index': i,
#                     'text': sentence
#                 })
#             if texts:
#                 # Use the underlying FAISS vectorstore
#                 self.vector_db.vectorstore.add_texts(texts, metadatas=metadatas)
#                 print(f"Added {len(texts)} sentences from {source} to vector DB.")

	# vectorstore.py
	import os
	import json
	from datetime import datetime
	# import faiss
	import numpy as np
	from sentence_transformers import SentenceTransformer
	from pypdf import PdfReader
	import requests
	from bs4 import BeautifulSoup

	# Updated LangChain imports
	from langchain_community.document_loaders import PyPDFLoader
	from langchain_community.document_loaders import WebBaseLoader
	from langchain_community.vectorstores import FAISS
	from langchain_huggingface import HuggingFaceEmbeddings


	class vectorstore:
	# 1 is pdf, 2 is website, and 3 is an already created vectorstore
	def __init__(self, path, Initlize_with=3):
	if Initlize_with == 1:
	self.index_path = self.create_vectorstore_from_pdf(path, embedder_model="all-MiniLM-L6-v2")
	self.vectorstore = self.load_vectorstore(self.index_path)
	elif Initlize_with == 2:
	self.index_path = self.create_vectorstore_from_website(path, embedder_model="all-MiniLM-L6-v2")
	self.vectorstore = self.load_vectorstore(self.index_path)
	elif Initlize_with == 3:
	self.index_path = path
	self.vectorstore = self.load_vectorstore(self.index_path)

	def chunk_text(self,text, chunk_size=1200):
	"""
	Split text into chunks of approximately 'chunk_size' words.

	Parameters:
	text (str): The input text.
	chunk_size (int): Maximum number of words per chunk.

	Returns:
	List[str]: A list of text chunks.
	"""
	words = text.split()
	chunks = []
	for i in range(0, len(words), chunk_size):
	chunk = " ".join(words[i:i+chunk_size])
	chunks.append(chunk)
	return chunks

	def create_vectorstore_from_pdf(self,pdf_path, chunk_size=1200, embedder_model="all-MiniLM-L6-v2"):
	"""
	Extract text from a PDF using pypdf, chunk the text, compute embeddings,
	and create a FAISS index.

	Parameters:
	pdf_path (str): Path to the PDF file.
	chunk_size (int): Number of words per chunk.
	embedder_model (str): The sentence-transformer model to use.

	Returns:
	index: A FAISS index containing the embeddings.
	chunks: A list of text chunks.
	embedder: The SentenceTransformer embedder.
	"""
	# Load PDF using LangChain's PyPDFLoader
	loader = PyPDFLoader(pdf_path)
	documents = loader.load()
	text = " ".join([doc.page_content for doc in documents])

	# Chunk text using existing method (word-based)
	chunks = self.chunk_text(text, chunk_size)

	# Create embeddings and FAISS vectorstore with LangChain
	embedder = HuggingFaceEmbeddings(model_name=embedder_model)
	vectorstore = FAISS.from_texts(chunks, embedder)

	# Save the vectorstore
	index_path = self.save_vectorstore_with_timestamp_and_without(vectorstore, embedder_model)
	return index_path

	def create_vectorstore_from_website(self, url, chunk_size=1200, embedder_model="all-MiniLM-L6-v2"):
	"""
	Fetch text from a website, chunk the text, compute embeddings,
	and create a FAISS index.

	Parameters:
	url (str): The URL of the website.
	chunk_size (int): Number of words per chunk.
	embedder_model (str): The sentence-transformer model to use.

	Returns:
	index: A FAISS index containing the embeddings.
	chunks: A list of text chunks.
	embedder: The SentenceTransformer embedder.
	"""
	# Load website using LangChain's WebBaseLoader
	loader = WebBaseLoader(url)
	documents = loader.load()
	text = " ".join([doc.page_content for doc in documents])

	# Chunk text using existing method
	chunks = self.chunk_text(text, chunk_size)

	# Create embeddings and FAISS vectorstore
	embedder = HuggingFaceEmbeddings(model_name=embedder_model)
	vectorstore = FAISS.from_texts(chunks, embedder)

	# Save the vectorstore
	index_path = self.save_vectorstore_with_timestamp_and_without(vectorstore, embedder_model)
	return index_path

	def add_to_vectorstore_web(self, url, index_file_path, chunk_size=1200):
	"""
	Fetch text from a website, chunk the text, compute embeddings using the existing embedder,
	and add them to the existing FAISS vectorstore.

	Parameters:
	url (str): The URL of the website.
	index_file_path (str): The file path of the saved vectorstore index.
	chunk_size (int): Maximum number of words per chunk.

	Returns:
	str: The updated vectorstore index file path.
	"""
	# Load the existing vectorstore
	vectorstore = self.load_vectorstore(index_file_path)

	# Fetch and extract text from the website
	loader = WebBaseLoader(url)
	documents = loader.load()
	text = " ".join([doc.page_content for doc in documents])

	# Chunk the text
	new_chunks = self.chunk_text(text, chunk_size)

	# Add new chunks to the vectorstore
	vectorstore.add_texts(new_chunks)

	# Save the updated vectorstore
	updated_index_path = self.save_vectorstore_with_timestamp_and_without(vectorstore, vectorstore.embedding_function.model_name)
	return updated_index_path


	def add_to_vectorstore_from_pdf(self, pdf_path, chunk_size=1200):
	"""
	Extract text from a PDF using PyPDF2, chunk the text, compute embeddings using the existing embedder,
	and add them to the existing FAISS vectorstore.

	Parameters:
	pdf_path (str): Path to the PDF file.
	index_file_path (str): The file path of the saved vectorstore index.
	chunk_size (int): Maximum number of words per chunk.

	Returns:
	str: The updated vectorstore index file path.
	"""
	# Extract text from PDF
	loader = PyPDFLoader(pdf_path)
	documents = loader.load()
	text = " ".join([doc.page_content for doc in documents])

	# Chunk the text
	new_chunks = self.chunk_text(text, chunk_size)

	# Add new chunks to the existing vectorstore
	self.vectorstore.add_texts(new_chunks)

	# Save the updated vectorstore
	self.index_path = self.save_vectorstore_with_timestamp_and_without(self.vectorstore, self.vectorstore.embedding_function.model_name)
	return self.index_path

	def save_vectorstore_with_timestamp_and_without(self, vectorstore, embedder_model=None):
	vector_db_folder = os.path.join(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")), "VectorDB")
	os.makedirs(vector_db_folder, exist_ok=True)

	timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
	folder_name = f"vectorstore_{timestamp}"
	folder_path = os.path.join(vector_db_folder, folder_name)

	# Save the LangChain FAISS vectorstore
	vectorstore.save_local(folder_path)

	# Save metadata with embedder model name
	metadata = {"embedder_model": embedder_model or "all-MiniLM-L6-v2"}
	with open(os.path.join(folder_path, "metadata.json"), "w") as f:
	json.dump(metadata, f)

	# Save to main folder
	main_folder = os.path.join(vector_db_folder, "vectorstore_mainV2")
	vectorstore.save_local(main_folder)
	with open(os.path.join(main_folder, "metadata.json"), "w") as f:
	json.dump(metadata, f)

	return main_folder

	def load_vectorstore(self, index_path):
	# Load metadata to get embedder model
	metadata_path = os.path.join(index_path, "metadata.json")
	with open(metadata_path, "r") as f:
	metadata = json.load(f)
	embedder_model = metadata["embedder_model"]

	# Create embedder and load vectorstore
	embedder = HuggingFaceEmbeddings(model_name=embedder_model)
	vectorstore = FAISS.load_local(index_path, embedder,allow_dangerous_deserialization=True)
	return vectorstore


	def search_vectorstore(self, query, top_k=5):
	# Search using LangChain's similarity_search
	docs = self.vectorstore.similarity_search(query, k=top_k)
	results = [doc.page_content for doc in docs]
	return results

	def log_conversation(self, user_text, bot_text=""):
	# Existing JSON logging
	logs_folder = os.path.join(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")), "logs")
	os.makedirs(logs_folder, exist_ok=True)
	log_file = os.path.join(logs_folder, "conversation_logs.json")
	entry = {
	"timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
	"User": {"text": user_text},
	"bot": {"text": bot_text}
	}
	if os.path.exists(log_file):
	with open(log_file, "r", encoding="utf-8") as f:
	try:
	logs = json.load(f)
	except json.JSONDecodeError:
	logs = []
	else:
	logs = []
	logs.append(entry)
	with open(log_file, "w", encoding="utf-8") as f:
	json.dump(logs, f, indent=4)

	# Add to LangMem memory
	self.memory.add_message({"role": "user", "content": user_text})
	if bot_text:
	self.memory.add_message({"role": "assistant", "content": bot_text})

	# # ...existing code...

	# import re
	# import requests
	# import nltk
	# from nltk.tokenize import sent_tokenize
	# import numpy as np

	# # Download NLTK data (only once; consider moving this to a setup step)
	# nltk.download('punkt')
	# nltk.download('punkt_tab')

	# class CancerDataIngestor:
	# def __init__(self,vector_db=None, source_urls=None, embedding_model=None):
	# """
	# Initialize with:
	# - source_urls: a dictionary mapping source names to URLs.
	# - embedding_model: a callable that takes text and returns an embedding.
	# - vector_db: an instance of your vector database.
	# """
	# self.sources = source_urls if source_urls is not None else {
	# 'PMC': 'https://www.ncbi.nlm.nih.gov/pmc/',
	# 'NCI': 'https://www.cancer.gov/about-cancer/understanding/statistics',
	# 'WHO': 'https://www.who.int/cancer/en/',
	# 'ClinicalTrials': 'https://clinicaltrials.gov/',
	# 'Kaggle': 'https://www.kaggle.com/datasets'
	# }
	# # Use the provided embedding model or default to self.get_embedding
	# self.embedding_model = embedding_model if embedding_model is not None else self.get_embedding
	# if vector_db is None:
	# raise ValueError("A valid vector DB instance must be provided.")
	# self.vector_db = vector_db
	# # Initialize the SentenceTransformer model once.
	# self.model = SentenceTransformer("all-MiniLM-L6-v2")

	# def fetch_html(self, url):
	# """Fetch HTML content from a given URL."""
	# try:
	# response = requests.get(url)
	# response.raise_for_status()
	# return response.text
	# except Exception as e:
	# print(f"Error fetching URL {url}: {e}")
	# return None

	# def clean_html(self, html):
	# """Extract text from HTML and remove tags and extra spaces."""
	# if html is None:
	# return ""
	# from bs4 import BeautifulSoup # in case not already imported above
	# soup = BeautifulSoup(html, 'html.parser')
	# for tag in soup(['script', 'style']):
	# tag.decompose()
	# text = soup.get_text(separator=' ')
	# text = re.sub(r'[^a-zA-Z0-9.,;:?!\s]', ' ', text)
	# text = re.sub(r'\s+', ' ', text)
	# return text.strip()

	# def preprocess_text(self, text):
	# """Lowercase the text and tokenize into sentences."""
	# text = text.lower()
	# sentences = sent_tokenize(text)
	# return sentences

	# def get_embedding(self, text):
	# """
	# Compute the embedding for a given text using the "all-MiniLM-L6-v2" model.
	# """
	# return self.model.encode(text)



	# def add_data_to_vectordb(self):
	# """
	# Iterate through all defined sources, fetch and clean the text,
	# generate embeddings for each sufficiently long sentence,
	# and add them to the vector DB using its underlying FAISS add_texts method.
	# """
	# for source, url in self.sources.items():
	# print(f"Processing source: {source}")
	# html_content = self.fetch_html(url)
	# cleaned_text = self.clean_html(html_content)
	# if not cleaned_text:
	# print(f"No text fetched from {url}")
	# continue
	# sentences = self.preprocess_text(cleaned_text)
	# texts = []
	# metadatas = []
	# for i, sentence in enumerate(sentences):
	# if len(sentence) < 50:
	# continue
	# texts.append(sentence)
	# metadatas.append({
	# 'source': source,
	# 'url': url,
	# 'sentence_index': i,
	# 'text': sentence
	# })
	# if texts:
	# # Use the underlying FAISS vectorstore
	# self.vector_db.vectorstore.add_texts(texts, metadatas=metadatas)
	# print(f"Added {len(texts)} sentences from {source} to vector DB.")