Upload folder using huggingface_hub

README.md

@@ -33,4 +33,27 @@ $ python main.py # or gradio main.py
 ```
 
 # Acknowledgments
-This is a project built during the Vertex sprints held by Google's ML Developer Programs team. We are thankful to be granted good amount of GCP credits to do this project.
+This is a project built during the Vertex sprints held by Google's ML Developer Programs team. We are thankful to be granted good amount of GCP credits to do this project.# AdaptSum
+
+AdaptSum stands for Adaptive Summarization. This project focuses on developing an LLM-powered system for dynamic summarization. Instead of generating entirely new summaries with each update, the system intelligently identifies and modifies only the necessary parts of the existing summary. This approach aims to create a more efficient and fluid summarization process within a continuous chat interaction with an LLM.
+
+# Instructions
+
+1. Install dependencies
+```shell
+$ pip install requirements.txt
+```
+
+2. Setup Gemini API Key
+```shell
+$ export GEMINI_API_KEY=xxxxx
+```
+> note that GEMINI API KEY should be obtained from Google AI Studio. Vertex AI is not supported at the moment (this is because Gemini SDK does not provide file uploading functionality for Vertex AI usage now).
+
+3. Run Gradio app
+```shell
+$ python main.py # or gradio main.py
+```
+
+# Acknowledgments
+This is a project built during the Vertex sprints held by Google's ML Developer Programs team. We are thankful to be granted good amount of GCP credits to do this project. 

configs/prompts.toml

@@ -15,15 +15,6 @@ When updating the summary:
 * **Completeness:**  Strive to be as comprehensive and detailed as possible, capturing all key points and insights.
 * **Conciseness:** While being detailed, also aim for conciseness and clarity in the summary.
 * **Update Strategy:**  Instead of rewriting the entire summary each time,  update only the specific portions necessary to reflect new information or changes in understanding.
-
-By following these guidelines, you will maintain an evolving summary that accurately reflects my learning and the key takeaways from our conversation."
-
-**Key improvements:**
-
-* **Clearer Instructions:** More explicit instructions on how to update the summary (i.e., updating specific portions instead of rewriting).
-* **Emphasis on Accuracy:** Stronger emphasis on factual accuracy and reflecting nuances.
-* **Conciseness:**  Added a direction to balance detail with conciseness.
-* **Structure:** Improved organization and formatting for better readability.
 """
 
 system_prompt = """
@@ -31,4 +22,4 @@ Consider yourself an expert at summarizing content with a high bar
 for scientific rigor. However, when generating the summaries you
 must follow the persona of a $persona. This persona will help set
 the tone of the conversation. 
-"""
+"""

hf_space/.github/workflows/sync_to_spaces.yml

@@ -0,0 +1,57 @@
+name: Sync to Hugging Face Spaces
+
+on:
+  push:
+    branches:
+      - main
+  workflow_dispatch:
+
+jobs:
+  sync:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout GitHub Repository
+        uses: actions/checkout@v3
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.10"
+
+      - name: Install Hugging Face Hub CLI
+        run: |
+          python -m pip install --upgrade pip
+          pip install huggingface_hub
+
+      - name: Clone Hugging Face Spaces Repository
+        run: |
+          huggingface-cli login --token $HF_TOKEN --add-to-git-credential
+          git clone https://huggingface.co/spaces/adaptsum/demo hf_space
+          cd hf_space
+          git checkout main
+        env:
+          HF_TOKEN: ${{ secrets.HUGGINGFACE_TOKEN }}
+
+      - name: Copy Files to Hugging Face Repo
+        run: |
+          rsync -av --exclude='.git' --exclude='README.md' ./ hf_space/
+
+      - name: Merge README.md Files
+        run: |
+          cat hf_space/README.md README.md > hf_space/README_combined.md
+          mv hf_space/README_combined.md hf_space/README.md
+          rm -rf hf_space/README_combined.md
+
+      - name: Commit and Push Changes
+        run: |
+          cd hf_space
+          git add .
+          if git diff --cached --quiet; then
+            echo "No changes to commit"
+          else
+            huggingface-cli upload adaptsum/demo . --repo-type=space
+            echo "Changes have been pushed."
+          fi
+        env:
+          HUGGINGFACE_TOKEN: ${{ secrets.HUGGINGFACE_TOKEN }}

hf_space/configs/prompts.toml

@@ -1,7 +1,7 @@
 [summarization]
 prompt = """
 Below is the initial summary of our conversation. 
-Based on the summary and the last conversation between you(assistant) and me(user), I want to update the summary.
+Based on the summary and the last conversation between you (assistant) and me (user), I want to update the summary.
 
 **Initial Summary:** 
 $previous_summary
@@ -23,5 +23,12 @@ By following these guidelines, you will maintain an evolving summary that accura
 * **Clearer Instructions:** More explicit instructions on how to update the summary (i.e., updating specific portions instead of rewriting).
 * **Emphasis on Accuracy:** Stronger emphasis on factual accuracy and reflecting nuances.
 * **Conciseness:**  Added a direction to balance detail with conciseness.
-* **Structure:** Improved organization and formatting for better readability. Bullet points with Markdown would be preferred.
+* **Structure:** Improved organization and formatting for better readability.
+"""
+
+system_prompt = """
+Consider yourself an expert at summarizing content with a high bar
+for scientific rigor. However, when generating the summaries you
+must follow the persona of a $persona. This persona will help set
+the tone of the conversation. 
 """

hf_space/hf_space/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

hf_space/hf_space/.gitignore

@@ -0,0 +1,171 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# PyPI configuration file
+.pypirc

hf_space/hf_space/app.py

@@ -0,0 +1,225 @@
+import os
+import argparse
+import asyncio
+import gradio as gr
+from difflib import Differ
+from string import Template
+from utils import load_prompt, setup_gemini_client
+from configs.responses import SummaryResponses
+from google.genai import types
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ai-studio-api-key", type=str, default=os.getenv("GEMINI_API_KEY"))
+    parser.add_argument("--vertexai", action="store_true", default=False)
+    parser.add_argument("--vertexai-project", type=str, default="gcp-ml-172005")
+    parser.add_argument("--vertexai-location", type=str, default="us-central1")
+    parser.add_argument("--model", type=str, default="gemini-2.0-flash", choices=["gemini-1.5-flash", "gemini-2.0-flash", "gemini-2.0-flash-001"])
+    parser.add_argument("--seed", type=int, default=2025)
+    parser.add_argument("--prompt-tmpl-path", type=str, default="configs/prompts.toml")
+    parser.add_argument("--css-path", type=str, default="statics/styles.css")
+    args = parser.parse_args()
+    return args
+
+def find_attached_file(filename, attached_files):
+    for file in attached_files:
+        if file['name'] == filename:
+            return file
+    return None
+
+async def echo(message, history, state, persona):
+    attached_file = None
+    system_instruction = Template(prompt_tmpl['summarization']['system_prompt']).safe_substitute(persona=persona)
+    
+    if message['files']:
+        path_local = message['files'][0]
+        filename = os.path.basename(path_local)
+
+        attached_file = find_attached_file(filename, state["attached_files"])
+        if attached_file is None: 
+            path_gcp = await client.files.upload(path=path_local)
+            state["attached_files"].append({
+                "name": filename,
+                "path_local": path_local,
+                "gcp_entity": path_gcp,
+                "path_gcp": path_gcp.name,
+                "mime_type=": path_gcp.mime_type,
+                "expiration_time": path_gcp.expiration_time,
+            })
+            attached_file = path_gcp
+
+    user_message = [message['text']]
+    if attached_file: user_message.append(attached_file)
+
+    chat_history = state['messages']
+    chat_history = chat_history + user_message
+    state['messages'] = chat_history
+
+    response_chunks = ""
+    model_content_stream = await client.models.generate_content_stream(
+    model=args.model, 
+    contents=state['messages'], 
+    config=types.GenerateContentConfig(
+        system_instruction=system_instruction, seed=args.seed
+    ),
+)
+    async for chunk in model_content_stream:
+        response_chunks += chunk.text
+        # when model generates too fast, Gradio does not respond that in real-time.
+        await asyncio.sleep(0.1)
+        yield (
+            response_chunks, 
+            state, 
+            state['summary_diff_history'][-1] if len(state['summary_diff_history']) > 1 else "",
+            state['summary_history'][-1] if len(state['summary_history']) > 1 else "",
+            gr.Slider(
+                visible=False if len(state['summary_history']) <= 1 else True, 
+                interactive=False if len(state['summary_history']) <= 1 else True, 
+            ),
+        )        
+    
+    # make summary
+    response = await client.models.generate_content(
+        model=args.model,
+        contents=[
+            Template(
+                prompt_tmpl['summarization']['prompt']
+            ).safe_substitute(
+                previous_summary=state['summary'], 
+                latest_conversation=str({"user": message['text'], "assistant": response_chunks})
+            )
+        ],
+        config=types.GenerateContentConfig(
+            system_instruction=system_instruction, 
+            seed=args.seed,
+            response_mime_type='application/json', 
+            response_schema=SummaryResponses
+        )
+    )
+
+    prev_summary = state['summary_history'][-1] if len(state['summary_history']) >= 1 else ""
+
+    state['summary'] = (
+        response.parsed.summary 
+        if getattr(response.parsed, "summary", None) is not None 
+        else response.text
+    )
+    state['summary_history'].append(
+        response.parsed.summary 
+        if getattr(response.parsed, "summary", None) is not None 
+        else response.text
+    )
+    state['summary_diff_history'].append(
+        [
+            (token[2:], token[0] if token[0] != " " else None)
+            for token in Differ().compare(prev_summary, state['summary'])
+        ]
+    )
+
+    yield (
+        response_chunks, 
+        state, 
+        state['summary_diff_history'][-1],
+        state['summary_history'][-1],
+        gr.Slider(
+            maximum=len(state['summary_history']),
+            value=len(state['summary_history']),
+            visible=False if len(state['summary_history']) == 1 else True, interactive=True
+        ),
+    )
+
+def change_view_toggle(view_toggle):
+    if view_toggle == "Diff":
+        return (
+            gr.HighlightedText(visible=True),
+            gr.Markdown(visible=False)
+        )
+    else:
+        return (
+            gr.HighlightedText(visible=False),
+            gr.Markdown(visible=True)
+        )        
+
+def navigate_to_summary(summary_num, state):
+    return (
+        state['summary_diff_history'][summary_num-1],
+        state['summary_history'][summary_num-1]
+    )
+
+def main(args):
+    style_css = open(args.css_path, "r").read()
+
+    global client, prompt_tmpl, system_instruction
+    client = setup_gemini_client(args)
+    prompt_tmpl = load_prompt(args)
+    
+    ## Gradio Blocks
+    with gr.Blocks(css=style_css) as demo:
+        # State per session
+        state = gr.State({
+            "messages": [],
+            "attached_files": [],
+            "summary": "",
+            "summary_history": [],
+            "summary_diff_history": []
+        })
+
+        with gr.Column():
+            gr.Markdown("# Adaptive Summarization")
+            gr.Markdown("AdaptSum stands for Adaptive Summarization. This project focuses on developing an LLM-powered system for dynamic summarization. Instead of generating entirely new summaries with each update, the system intelligently identifies and modifies only the necessary parts of the existing summary. This approach aims to create a more efficient and fluid summarization process within a continuous chat interaction with an LLM.")
+
+        with gr.Column():
+            with gr.Accordion("Adaptively Summarized Conversation", elem_id="adaptive-summary-accordion", open=False):
+                with gr.Row(elem_id="view-toggle-btn-container"):
+                    view_toggle_btn = gr.Radio(
+                        choices=["Diff", "Markdown"],
+                        value="Markdown",
+                        interactive=True,
+                        elem_id="view-toggle-btn"
+                    )
+
+                summary_diff = gr.HighlightedText(
+                    label="Summary so far",
+                    # value="No summary yet. As you chat with the assistant, the summary will be updated automatically.",
+                    combine_adjacent=True,
+                    show_legend=True,
+                    color_map={"-": "red", "+": "green"},
+                    elem_classes=["summary-window"],
+                    visible=False
+                )
+
+                summary_md = gr.Markdown(
+                    label="Summary so far",
+                    value="No summary yet. As you chat with the assistant, the summary will be updated automatically.",
+                    elem_classes=["summary-window"],
+                    visible=True
+                )
+
+                summary_num = gr.Slider(label="summary history", minimum=1, maximum=1, step=1, show_reset_button=False, visible=False)
+
+            view_toggle_btn.change(change_view_toggle, inputs=[view_toggle_btn], outputs=[summary_diff, summary_md])
+            summary_num.release(navigate_to_summary, inputs=[summary_num, state], outputs=[summary_diff, summary_md])
+        
+        with gr.Column("persona-dropdown-container", elem_id="persona-dropdown-container"):
+            persona = gr.Dropdown(
+                ["expert", "novice", "regular practitioner", "high schooler"], 
+                label="Summary Persona", 
+                info="Control the tonality of the conversation.",
+                min_width="auto",
+            )        
+
+        with gr.Column("chat-window", elem_id="chat-window"):
+            gr.ChatInterface(
+                multimodal=True,
+                type="messages", 
+                fn=echo, 
+                additional_inputs=[state, persona],
+                additional_outputs=[state, summary_diff, summary_md, summary_num],
+            )
+
+    return demo
+
+if __name__ == "__main__":
+    args = parse_args()
+    demo = main(args)
+    demo.launch()

hf_space/hf_space/configs/prompts.toml

@@ -0,0 +1,27 @@
+[summarization]
+prompt = """
+Below is the initial summary of our conversation. 
+Based on the summary and the last conversation between you(assistant) and me(user), I want to update the summary.
+
+**Initial Summary:** 
+$previous_summary
+
+**Last Conversation:**
+$latest_conversation
+
+When updating the summary:
+* **Focus:**  Only include information we have explicitly discussed in this session. Do not introduce any new information or topics, even if you have prior knowledge.
+* **Accuracy:** Ensure the summary is factually accurate and reflects the nuances of our discussion.
+* **Completeness:**  Strive to be as comprehensive and detailed as possible, capturing all key points and insights.
+* **Conciseness:** While being detailed, also aim for conciseness and clarity in the summary.
+* **Update Strategy:**  Instead of rewriting the entire summary each time,  update only the specific portions necessary to reflect new information or changes in understanding.
+
+By following these guidelines, you will maintain an evolving summary that accurately reflects my learning and the key takeaways from our conversation."
+
+**Key improvements:**
+
+* **Clearer Instructions:** More explicit instructions on how to update the summary (i.e., updating specific portions instead of rewriting).
+* **Emphasis on Accuracy:** Stronger emphasis on factual accuracy and reflecting nuances.
+* **Conciseness:**  Added a direction to balance detail with conciseness.
+* **Structure:** Improved organization and formatting for better readability. Bullet points with Markdown would be preferred.
+"""

hf_space/hf_space/configs/responses.py

@@ -0,0 +1,4 @@
+from pydantic import BaseModel
+
+class SummaryResponses(BaseModel):
+    summary: str

hf_space/hf_space/requirements.txt

@@ -0,0 +1,3 @@
+google-genai==0.7.0
+toml
+gradio

hf_space/hf_space/statics/styles.css

@@ -0,0 +1,54 @@
+.summary-window {
+    height: 550px !important;
+    /* border: dashed 1px #e0e0e0 !important; */
+    border-radius: 10px !important;
+    padding: 4px;
+}
+
+.summary-window > label {
+    display: none !important;
+}
+
+#view-toggle-btn-container > div {
+    border: none !important;
+}
+
+#view-toggle-btn > span {
+    display: none !important;
+}
+
+#view-toggle-btn > div:nth-child(3) {
+    margin: auto !important;
+    width: fit-content !important;
+}
+
+#adaptive-summary-accordion {
+    position: absolute !important;
+    z-index: 100 !important;
+    box-shadow: 5px 5px 10px rgba(0, 0, 0, 0.2);
+}
+
+#chat-window {
+    margin-top: 40px !important;
+}
+
+#chat-window > div > div:nth-child(1) {
+    height: 600px !important;
+}
+
+.textfield {
+    line-height: 1.7 !important;
+}
+
+.textspan {
+    padding: 0px !important;
+    margin: 0px !important;
+    line-height: 1.7 !important;
+}
+
+@media (prefers-color-scheme: dark) {
+    #adaptive-summary-accordion {
+        /* White-ish shadow for dark themes */
+        box-shadow: 5px 5px 10px rgba(245, 245, 245, 0.2); /* Or any other white-ish color you prefer */
+    }
+}

hf_space/hf_space/utils.py

@@ -0,0 +1,22 @@
+import toml
+from google import genai
+
+def load_prompt(args):
+    with open(args.prompt_tmpl_path, 'r') as f:
+        prompts = toml.load(f)
+
+    return prompts
+
+def setup_gemini_client(args):
+    if args.vertexai:
+        client = genai.Client(
+            vertexai=args.vertexai, 
+            project=args.vertexai_project, 
+            location=args.vertexai_location
+        )
+    else:
+        client = genai.Client(
+            api_key=args.ai_studio_api_key,
+        )
+
+    return client

hf_space/requirements.txt

@@ -1,3 +1,3 @@
 google-genai==0.7.0
 toml
-gradio
+gradio

hf_space/statics/styles.css

@@ -1,5 +1,5 @@
 .summary-window {
-    height: 550px !important;
+    height: 600px !important;
     /* border: dashed 1px #e0e0e0 !important; */
     border-radius: 10px !important;
     padding: 4px;
@@ -28,9 +28,9 @@
     box-shadow: 5px 5px 10px rgba(0, 0, 0, 0.2);
 }
 
-#chat-window {
+/* #chat-window {
     margin-top: 40px !important;
-}
+} */
 
 #chat-window > div > div:nth-child(1) {
     height: 600px !important;
@@ -46,9 +46,6 @@
     line-height: 1.7 !important;
 }
 
-@media (prefers-color-scheme: dark) {
-    #adaptive-summary-accordion {
-        /* White-ish shadow for dark themes */
-        box-shadow: 5px 5px 10px rgba(245, 245, 245, 0.2); /* Or any other white-ish color you prefer */
-    }
+#persona-dropdown-container {
+    margin-top: 40px !important;
 }

hf_space/utils.py

@@ -9,14 +9,18 @@ def load_prompt(args):
 
 def setup_gemini_client(args):
     if args.vertexai:
-        client = genai.Client(
-            vertexai=args.vertexai, 
-            project=args.vertexai_project, 
-            location=args.vertexai_location
+        client = genai.client.AsyncClient(
+            genai.client.ApiClient(
+                vertexai=args.vertexai, 
+                project=args.vertexai_project, 
+                location=args.vertexai_location
+            )
         )
     else:
-        client = genai.Client(
-            api_key=args.ai_studio_api_key,
+        client = genai.client.AsyncClient(
+            genai.client.ApiClient(
+                api_key=args.ai_studio_api_key
+            )
         )
 
     return client