Upload main.py

main.py

@@ -0,0 +1,121 @@
+import os
+import fitz  # PyMuPDF
+from docx import Document
+from sentence_transformers import SentenceTransformer
+import faiss
+import numpy as np
+import pickle
+from langchain_community.llms import HuggingFaceEndpoint
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+
+# Function to extract text from a PDF file
+def extract_text_from_pdf(pdf_path):
+    text = ""
+    doc = fitz.open(pdf_path)
+    for page_num in range(len(doc)):
+        page = doc.load_page(page_num)
+        text += page.get_text()
+    return text
+
+# Function to extract text from a Word document
+def extract_text_from_docx(docx_path):
+    doc = Document(docx_path)
+    text = "\n".join([para.text for para in doc.paragraphs])
+    return text
+
+# Initialize the embedding model
+embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+
+# Path to the document (can be either a single file or a directory)
+docs_path = "C:\\Users\\MOD\\chatbot\\Should companies implement a four.docx"
+
+documents = []
+doc_texts = []
+
+if os.path.isdir(docs_path):
+    # Iterate through all files in the directory
+    for filename in os.listdir(docs_path):
+        file_path = os.path.join(docs_path, filename)
+        if filename.endswith(".pdf"):
+            text = extract_text_from_pdf(file_path)
+            documents.append(filename)
+            doc_texts.append(text)
+        elif filename.endswith(".docx"):
+            text = extract_text_from_docx(file_path)
+            documents.append(filename)
+            doc_texts.append(text)
+elif os.path.isfile(docs_path):
+    # Process a single file
+    if docs_path.endswith(".pdf"):
+        text = extract_text_from_pdf(docs_path)
+        documents.append(os.path.basename(docs_path))
+        doc_texts.append(text)
+    elif docs_path.endswith(".docx"):
+        text = extract_text_from_docx(docs_path)
+        documents.append(os.path.basename(docs_path))
+        doc_texts.append(text)
+else:
+    print("Invalid path specified. Please provide a valid file or directory path.")
+
+# Generate embeddings for the document texts
+embeddings = embedding_model.encode(doc_texts)
+
+# Create a FAISS index
+d = embeddings.shape[1]  # Dimension of the embeddings
+index = faiss.IndexFlatL2(d)  # L2 distance metric
+index.add(np.array(embeddings))  # Add embeddings to the index
+
+# Save the FAISS index and metadata
+index_path = "faiss_index"
+if not os.path.exists(index_path):
+    os.makedirs(index_path)
+
+faiss.write_index(index, os.path.join(index_path, "index.faiss"))
+
+# Save the document metadata to a file for retrieval purposes
+with open(os.path.join(index_path, "documents.txt"), "w") as f:
+    for doc in documents:
+        f.write("%s\n" % doc)
+
+# Save additional metadata
+metadata = {
+    "documents": documents,
+    "embeddings": embeddings
+}
+with open(os.path.join(index_path, "index.pkl"), "wb") as f:
+    pickle.dump(metadata, f)
+
+print("FAISS index and documents saved.")
+
+# Load the FAISS index and metadata
+index = faiss.read_index(os.path.join(index_path, "index.faiss"))
+with open(os.path.join(index_path, "index.pkl"), "rb") as f:
+    metadata = pickle.load(f)
+documents = metadata["documents"]
+embeddings = metadata["embeddings"]
+
+# Retrieve the API token from the environment variable
+api_token = os.getenv('HUGGINGFACEHUB_API_TOKEN')
+if api_token is None:
+    raise ValueError("HUGGINGFACEHUB_API_TOKEN environment variable is not set")
+
+print(f"API Token: {api_token[:5]}...")  # Print the first 5 characters of the token for verification
+
+# Initialize the LLM
+llm = HuggingFaceEndpoint(
+    endpoint_url="https://api-inference.huggingface.co/models/gpt2",
+    model_kwargs={"api_key": api_token}
+)
+
+# Function to perform a search query
+def search(query, k=5):
+    query_embedding = embedding_model.encode([query])
+    D, I = index.search(np.array(query_embedding), k)
+    results = [documents[i] for i in I[0]]
+    return results
+
+# Example query
+query = "What is the impact of a four-day work week?"
+results = search(query)
+print("Top documents:", results)