Merge branch 'main' into feature/test-gitflow

zenml-io · Nov 25, 2024 · e851e4a · e851e4a
2 parents 33fc151 + cf029d1
commit e851e4a
Show file tree

Hide file tree

Showing 15 changed files with 337 additions and 74 deletions.
diff --git a/llm-complete-guide/ZENML_VERSION.txt b/llm-complete-guide/ZENML_VERSION.txt
@@ -1 +1 @@
-0.71.0
+0.71.0
diff --git a/llm-complete-guide/configs/dev/rag.yaml b/llm-complete-guide/configs/dev/rag.yaml
@@ -16,6 +16,9 @@ settings:
       - rerankers
       - pygithub
       - rerankers[flashrank]
+      - matplotlib
+      - elasticsearch
+
     environment:
       ZENML_PROJECT_SECRET_NAME: llm_complete
       ZENML_ENABLE_RICH_TRACEBACK: FALSE

diff --git a/llm-complete-guide/configs/dev/rag_eval.yaml b/llm-complete-guide/configs/dev/rag_eval.yaml
@@ -13,4 +13,5 @@ settings:
       - psycopg2-binary
       - tiktoken
       - pygithub
+      - elasticsearch
     python_package_installer: "uv"
diff --git a/llm-complete-guide/configs/production/eval.yaml b/llm-complete-guide/configs/production/eval.yaml
@@ -17,6 +17,7 @@ settings:
       - matplotlib
       - pillow
       - pygithub
+      - elasticsearch
     environment:
       ZENML_PROJECT_SECRET_NAME: llm_complete
       ZENML_ENABLE_RICH_TRACEBACK: FALSE

diff --git a/llm-complete-guide/configs/production/rag.yaml b/llm-complete-guide/configs/production/rag.yaml
@@ -1,4 +1,4 @@
-enable_cache: True
+enable_cache: False
 
 # environment configuration
 settings:
@@ -16,6 +16,9 @@ settings:
       - rerankers
       - pygithub
       - rerankers[flashrank]
+      - matplotlib
+      - elasticsearch
+
     environment:
       ZENML_PROJECT_SECRET_NAME: llm_complete
       ZENML_ENABLE_RICH_TRACEBACK: FALSE
@@ -26,7 +29,6 @@ steps:
     parameters:
       docs_url: https://docs.zenml.io
       use_dev_set: false
-    enable_cache: true
 #  generate_embeddings:
 #    step_operator: "sagemaker"
 #    settings:

diff --git a/llm-complete-guide/configs/staging/embeddings.yaml b/llm-complete-guide/configs/staging/embeddings.yaml
@@ -1,4 +1,4 @@
-# enable_cache: False
+enable_cache: False
 
 # environment configuration
 settings:

diff --git a/llm-complete-guide/configs/staging/eval.yaml b/llm-complete-guide/configs/staging/eval.yaml
@@ -17,6 +17,7 @@ settings:
       - matplotlib
       - pillow
       - pygithub
+      - elasticsearch
     environment:
       ZENML_PROJECT_SECRET_NAME: llm_complete
       ZENML_ENABLE_RICH_TRACEBACK: FALSE

diff --git a/llm-complete-guide/configs/staging/rag.yaml b/llm-complete-guide/configs/staging/rag.yaml
@@ -16,6 +16,9 @@ settings:
       - rerankers
       - pygithub
       - rerankers[flashrank]
+      - matplotlib
+      - elasticsearch
+
     environment:
       ZENML_PROJECT_SECRET_NAME: llm_complete
       ZENML_ENABLE_RICH_TRACEBACK: FALSE
@@ -29,4 +32,3 @@ steps:
     parameters:
       docs_url: https://docs.zenml.io
       use_dev_set: false
-    enable_cache: true
diff --git a/llm-complete-guide/constants.py b/llm-complete-guide/constants.py
@@ -23,6 +23,9 @@
     384  # Update this to match the dimensionality of the new model
 )
 
+# ZenML constants
+ZENML_CHATBOT_MODEL = "zenml-docs-qa-chatbot"
+
 # Scraping constants
 RATE_LIMIT = 5  # Maximum number of requests per second
 
@@ -78,3 +81,4 @@
 USE_ARGILLA_ANNOTATIONS = False
 
 SECRET_NAME = os.getenv("ZENML_PROJECT_SECRET_NAME", "llm-complete")
+SECRET_NAME_ELASTICSEARCH = "elasticsearch-zenml"
diff --git a/llm-complete-guide/gh_action_rag.py b/llm-complete-guide/gh_action_rag.py
@@ -101,8 +101,6 @@ def main(
         zenml_model_name (str): The ZenML model name.
         zenml_model_version (str): The ZenML model version.
     """
-    pipeline_args = {"enable_cache": not no_cache}
-
     client = Client()
     config_path = Path(__file__).parent / "configs" / config
 

diff --git a/llm-complete-guide/pipelines/finetune_embeddings.py b/llm-complete-guide/pipelines/finetune_embeddings.py
@@ -20,16 +20,10 @@
     prepare_load_data,
     visualize_results,
 )
-from zenml import Model, pipeline
-from zenml.model.model import ModelStages
-
-model_definition = Model(
-    name=EMBEDDINGS_MODEL_NAME_ZENML,
-    version=ModelStages.LATEST,
-)
+from zenml import pipeline
 
 
-@pipeline(model=model_definition)
+@pipeline
 def finetune_embeddings():
     data = prepare_load_data()
     base_results = evaluate_base_model(dataset=data)

diff --git a/llm-complete-guide/requirements.txt b/llm-complete-guide/requirements.txt
@@ -1,4 +1,4 @@
-zenml[server]>=0.68.1
+zenml[server]==0.68.1
 ratelimit
 pgvector
 psycopg2-binary
@@ -20,6 +20,7 @@ datasets
 torch
 gradio
 huggingface-hub
+elasticsearch
 
 # optional requirements for S3 artifact store
 # s3fs>2022.3.0

diff --git a/llm-complete-guide/steps/eval_retrieval.py b/llm-complete-guide/steps/eval_retrieval.py
@@ -19,8 +19,10 @@
 
 from datasets import load_dataset
 from utils.llm_utils import (
+    find_vectorstore_name,
     get_db_conn,
     get_embeddings,
+    get_es_client,
     get_topn_similar_docs,
     rerank_documents,
 )
@@ -76,11 +78,23 @@ def query_similar_docs(
         Tuple containing the question, URL ending, and retrieved URLs.
     """
     embedded_question = get_embeddings(question)
-    db_conn = get_db_conn()
+    conn = None
+    es_client = None
+
+    vector_store_name = find_vectorstore_name()
+    if vector_store_name == "pgvector":
+        conn = get_db_conn()
+    else:
+        es_client = get_es_client()
+
     num_docs = 20 if use_reranking else returned_sample_size
     # get (content, url) tuples for the top n similar documents
     top_similar_docs = get_topn_similar_docs(
-        embedded_question, db_conn, n=num_docs, include_metadata=True
+        embedded_question, 
+        conn=conn, 
+        es_client=es_client, 
+        n=num_docs, 
+        include_metadata=True
     )
 
     if use_reranking: