add gemini token counter

app.py

@@ -14,6 +14,24 @@ from streamlit_pdf_viewer import pdf_viewer
 
 MAX_OUTPUT_TOKENS = 2048  
 
+
+class CountGeminiTokens:
+    """
+    Count tokens in Gemini models.
+    
+    See: https://medium.com/google-cloud/counting-gemini-text-tokens-locally-with-the-vertex-ai-sdk-78979fea6244
+    """
+    def __init__(self, llm_name):
+        from vertexai.preview import tokenization
+        self.tokenizer = tokenization.get_tokenizer_for_model(llm_name)
+
+    def __call__(self, input):
+        """This returns all the tokens in a list since LlamaIndex seems to count by calling `len()` on the tokenizer function."""
+        tokens = []
+        for list in self.tokenizer.compute_tokens(input).token_info_list:
+            tokens += list.tokens
+        return tokens
+        
 def main():
     with st.sidebar:
         st.title('Document Summarization and QA System')
@@ -84,7 +102,6 @@ def main():
                 # raise NotImplementedError(f"{provider} is not supported yet")
                 from llama_index.llms.gemini import Gemini
                 from llama_index.embeddings.gemini import GeminiEmbedding
-                from vertexai.preview import tokenization
 
                 os.environ['GOOGLE_API_KEY'] = str(llm_key)
                 Settings.llm = Gemini(
@@ -93,7 +110,7 @@ def main():
                     temperature=temperature,
                     max_tokens=MAX_OUTPUT_TOKENS
                 )
-                Settings.tokenizer = tokenization.get_tokenizer_for_model(llm_name).compute_tokens
+                Settings.tokenizer = CountGeminiTokens(llm_name) #tokenization.get_tokenizer_for_model(llm_name).compute_tokens
                 Settings.num_output = MAX_OUTPUT_TOKENS
                 Settings.embed_model = GeminiEmbedding(
                     model_name="models/text-embedding-004", api_key=os.environ.get("GOOGLE_API_KEY") #, title="this is a document"