minor fixes

app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import os
 import pandas as pd
 from command_center import CommandCenter
-from process_documents import process_documents
+from process_documents import process_documents, num_tokens
 from embed_documents import create_retriever
 import json
 from langchain.callbacks import get_openai_callback
@@ -59,7 +59,12 @@ def process_documents_wrapper(inputs):
     st.session_state.retriever = create_retriever(snippets)
     st.session_state.source_doc_urls = inputs
     st.session_state.index = [
-        [snip.metadata["chunk_id"], snip.metadata["header"]] for snip in snippets
+        [
+            snip.metadata["chunk_id"],
+            snip.metadata["header"],
+            num_tokens(snip.page_content),
+        ]
+        for snip in snippets
     ]
     response = f"Uploaded and processed documents {inputs}"
     st.session_state.messages.append((f"/add-papers {inputs}", response, "identity"))
@@ -68,7 +73,9 @@ def process_documents_wrapper(inputs):
 
 
 def index_documents_wrapper(inputs=None):
-    response = pd.DataFrame(st.session_state.index, columns=["id", "reference"])
+    response = pd.DataFrame(
+        st.session_state.index, columns=["id", "reference", "tokens"]
+    )
     st.session_state.messages.append(("/library", response, "dataframe"))
     return (response, "dataframe")
 

process_documents.py

@@ -4,6 +4,7 @@ from statistics import median
 from bs4 import BeautifulSoup
 from langchain.docstore.document import Document
 from langchain.document_loaders import PDFMinerPDFasHTMLLoader, WebBaseLoader
+import tiktoken
 
 deep_strip = lambda text: re.sub(r"\s+", " ", text or "").strip()
 
@@ -153,3 +154,7 @@ def get_pdf_semantic_snippets(filtered_snippets, median_font_size):
         }
         semantic_snippets.append((current_content, metadata))
     return semantic_snippets
+
+
+def num_tokens(string):
+    return len(tiktoken.get_encoding("cl100k_base").encode(string))