Create ingest.py

ingest.py

@@ -0,0 +1,37 @@
+import os
+import torch
+from constants import CHROMA_SETTINGS
+from langchain.document_loaders import PDFMinerLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores import Chroma
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+checkpoint = "MBZUAI/LaMini-T5-738M"
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint, device_map="auto", torch_dtype=torch.float32)
+
+persist_directory = "db"
+
+def main():
+    for root, dirs, files in os.walk("docs"):
+        for file in files:
+            if file.endswith(".pdf"):
+                print(f"Ingesting file: {file}")
+                loader = PDFMinerLoader(os.path.join(root, file))
+                documents = loader.load()
+                
+                text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+                texts = text_splitter.split_documents(documents)
+
+                def embedding_function(text):
+                    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(model.device)
+                    with torch.no_grad():
+                        embeddings = model.encoder(**inputs).last_hidden_state.mean(dim=1).cpu().numpy()
+                    return embeddings
+
+                db = Chroma.from_documents(texts, embedding_function=embedding_function, persist_directory=persist_directory, client_settings=CHROMA_SETTINGS)
+                db.persist()
+                db = None
+
+if __name__ == "__main__":
+    main()