File size: 2,063 Bytes
4f0811d
 
 
 
 
 
72e83b2
4f0811d
 
dca455a
4f0811d
dca455a
 
4f0811d
dca455a
4f0811d
dca455a
 
 
4f0811d
dca455a
4f0811d
dca455a
4f0811d
dca455a
b329506
4f0811d
 
dca455a
 
 
4f0811d
dca455a
 
4f0811d
 
 
dca455a
4f0811d
 
dca455a
 
4f0811d
dca455a
 
4f0811d
dca455a
4f0811d
dca455a
 
 
 
4f0811d
dca455a
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
import os
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.document_loaders import PyMuPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from langchain.llms import HuggingFaceHub

# Constants
CHROMA_DB_PATH = "chroma_db"
SENTENCE_TRANSFORMER_MODEL = "sentence-transformers/all-MiniLM-L6-v2"  # Corrected model name
LLM_MODEL = "HuggingFaceH4/zephyr-7b-beta"  # Free chatbot model from Hugging Face

# Initialize vector store
def initialize_vector_store():
    """Initialize or load ChromaDB vector store"""
    embeddings = HuggingFaceEmbeddings(model_name=SENTENCE_TRANSFORMER_MODEL)
    vector_store = Chroma(persist_directory=CHROMA_DB_PATH, embedding_function=embeddings)
    return vector_store

vector_store = initialize_vector_store()

def ingest_pdf(pdf_path):
    """Processes a PDF, splits text, and stores embeddings in ChromaDB."""
    loader = PyMuPDFLoader(pdf_path)
    documents = loader.load()

    # Split text into smaller chunks
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
    split_docs = text_splitter.split_documents(documents)

    # Store in vector database
    vector_store.add_documents(split_docs)
    vector_store.persist()

def process_query_with_memory(query, chat_history=[]):
    """Retrieves relevant document chunks and generates a conversational response."""
    retriever = vector_store.as_retriever()

    # Initialize chat memory
    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

    # Load a free Hugging Face model
    llm = HuggingFaceHub(repo_id=LLM_MODEL, model_kwargs={"max_new_tokens": 500})

    # Create a conversational retrieval chain
    qa_chain = ConversationalRetrievalChain(
        llm=llm,
        retriever=retriever,
        memory=memory
    )

    return qa_chain.run({"question": query, "chat_history": chat_history})