hacker mode

add_data.py

@@ -2,6 +2,19 @@ import chromadb
 from chromadb.utils import embedding_functions
 
 
+def create_client():
+    client = chromadb.PersistentClient(path="./chromadb_linux/")
+    MODEL_NAME: str = "mixedbread-ai/mxbai-embed-large-v1"  # ~ 0.5 gb
+    COLLECTION_NAME: str = "schemer2"
+    EMBEDDING_FUNC = embedding_functions.SentenceTransformerEmbeddingFunction(
+        model_name=MODEL_NAME
+    )
+    schemer = client.get_collection(
+        name=COLLECTION_NAME,
+        embedding_function=EMBEDDING_FUNC,
+    )
+    return schemer
+
 def get_client():
     client = chromadb.PersistentClient(path="./chromadb_linux/")
     MODEL_NAME: str = "mixedbread-ai/mxbai-embed-large-v1"  # ~ 0.5 gb
@@ -20,11 +33,14 @@ def update_collection(iter: int, text: object, client: chromadb.Collection):
     client.add(documents=[text["text"]], metadatas=[{"source": "pdf"}], ids=[text["content"] + str(iter)])
 
 
-def encode_image(img_path: str):
+def encode_image(image) -> str:
+    import io
     import base64
-
-    with open(img_path, "rb") as image_file:
-        return base64.b64encode(image_file.read()).decode("utf-8")
+    
+    byte_arr = io.BytesIO()
+    image.save(byte_arr, format="JPEG")
+    encoded_image = base64.b64encode(byte_arr.getvalue()).decode("utf-8")
+    return encoded_image
 
 
 async def image_to_text(image) -> object:
@@ -55,28 +71,26 @@ async def image_to_text(image) -> object:
     return json.loads(response.choices[0].message.content)
 
 
-async def start_troggin_off(dir: str):
+async def start_troggin_off(dir: str, client):
+    # recursive
     import os
     from pdf2image import convert_from_path
 
-    client = get_client()
-
-    for folder in os.listdir(dir):
-        folder_path = os.path.join(dir, folder)
-        if os.path.isdir(folder_path):
-            for file in os.listdir(folder_path):
-                if file.endswith(".pdf"):
-                    print("Processing", file)
-                    pdf_path = os.path.join(folder_path, file)
-                    images = convert_from_path(pdf_path)
+    dirs = os.listdir(dir)
+    for path in dirs:
+        if os.path.isdir(os.path.join(dir, path)):
+            await start_troggin_off(os.path.join(dir, path), client) # recursive call
 
-                    for i, image in enumerate(images):
-                        image.save(f"out{i}.jpg", "JPEG")
-                        encoded_image = encode_image(f"out{i}.jpg")
-                        text = await image_to_text(encoded_image)
-                        update_collection(i, text, client)
+        if(os.path.join(dir, path).endswith(".pdf")):
+            images = convert_from_path(os.path.join(dir, path))
 
+            for i, image in enumerate(images):
+                encoded_image = encode_image(image)
+                text = await image_to_text(encoded_image)
+                update_collection(i, text, client)
 
 if __name__ == "__main__":
     import asyncio
-    asyncio.run(start_troggin_off("data/Class Notes/"))
+    client = create_client()
+    # client = None
+    asyncio.run(start_troggin_off("data/", client))

app.py

@@ -2,8 +2,10 @@ import streamlit as st
 from openai import OpenAI
 from params import params
 from database import get_client
+from add_data import start_troggin_off, create_client
 
-CLIENT = get_client()
+# CLIENT = get_client()
+CLIENT = None
 APP_NAME: str = "Groove-GPT"
 history = []
 st.set_page_config(layout="wide")
@@ -11,6 +13,12 @@ st.set_page_config(layout="wide")
 # INFO
 st.title(APP_NAME)
 
+
+start_embedding = st.button("Hacker man")
+if start_embedding:
+    CLIENT = create_client()
+    start_troggin_off("./data", CLIENT)
+
 l_col, r_col = st.columns((3, 1))
 
 # param column
@@ -41,7 +49,7 @@ with l_col:
             )
             documents = results["documents"]
             response = openai_client.chat.completions.create(
-                model="gpt-3.5-turbo",
+                model=gpt_type,
                 messages=[
                     {
                         "role": "system",