Add driver proxy generator (databrickslabs#10)

suanmiao · web-flow · commit 3780ee38fcef · 2023-10-02T15:39:23.000-07:00
diff --git a/databricks/labs/doc_qa/chatbot/chatbot.py b/databricks/labs/doc_qa/chatbot/chatbot.py
@@ -2,11 +2,8 @@
 from databricks.labs.doc_qa.llm_utils import PromptTemplate
 import openai
 from databricks.labs.doc_qa.chatbot.retriever import Document, BaseRetriever
-import logging
+from databricks.labs.doc_qa.logging_utils import logger
 import tiktoken
-import dataclasses
-
-logger = logging.getLogger(__name__)
 
 class LlmProvider:
 
@@ -16,9 +13,9 @@ def prompt(self, prompt: str, **kwargs) -> str:
 
 class OpenAILlmProvider(LlmProvider):
 
-    def __init__(self, api_key: str, model, temperature, **kwargs):
+    def __init__(self, api_key: str, model_name: str, temperature, **kwargs):
         self._api_key = api_key
-        self._model = model
+        self.model_name = model_name
         self._temperature = temperature
         openai.api_key = api_key
     
@@ -51,7 +48,7 @@ def __init__(self, llm_provider: str, retriever: BaseRetriever,
         self._whole_prompt_template = whole_prompt_template
         self._document_prompt_tempate = document_prompt_tempate
         self._max_num_tokens_for_context = max_num_tokens_for_context
-        self._enc = tiktoken.encoding_for_model(self._llm_provider.model)
+        self._enc = tiktoken.encoding_for_model(self._llm_provider.model_name)
 
     def chat(self, prompt: str, top_k=1, **kwargs) -> ChatResponse:
         """
diff --git a/databricks/labs/doc_qa/chatbot/retriever.py b/databricks/labs/doc_qa/chatbot/retriever.py
@@ -1,15 +1,13 @@
 from datetime import datetime
 import pandas as pd
 from databricks.labs.doc_qa.llm_utils import PromptTemplate
+from databricks.labs.doc_qa.logging_utils import logger
 import openai
-import logging
 import faiss
 import numpy as np
 import json
 
 
-logger = logging.getLogger(__name__)
-
 
 class EmbeddingProvider:
 
diff --git a/databricks/labs/doc_qa/doc_qa_helper.py b/databricks/labs/doc_qa/doc_qa_helper.py
@@ -0,0 +1,113 @@
+from databricks.labs.doc_qa.llm_utils import PromptTemplate
+import pandas as pd
+import os
+from databricks.labs.doc_qa.evaluators.templated_evaluator import (
+    OpenAIEvaluator,
+    RetryPolicy,
+)
+from databricks.labs.doc_qa.variables.doc_qa_template_variables import (
+    get_openai_grading_template_and_function,
+)
+from databricks.labs.doc_qa.logging_utils import logger
+
+
+def gpt_4_evaluator():
+    retry_policy = RetryPolicy(max_retry_on_invalid_result=3, max_retry_on_exception=3)
+    (
+        openai_grading_prompt,
+        openai_grading_function,
+    ) = get_openai_grading_template_and_function(scale=10, level_of_details=2)
+    return OpenAIEvaluator(
+        model="gpt-4",
+        temperature=0.0,
+        grading_prompt_tempate=openai_grading_prompt,
+        input_columns=["question", "answer", "context"],
+        openai_function=openai_grading_function,
+        retry_policy=retry_policy,
+    )
+
+
+def vllm_vicuna_model_generator(url, pat_token, model_name):
+    from databricks.labs.doc_qa.model_generators.model_generator import (
+        vLllmOpenAICompletionFormatModelGenerator,
+    )
+    from databricks.labs.doc_qa.variables.doc_qa_template_variables import (
+        vicuna_prompt_format_func,
+        doc_qa_task_prompt_template,
+    )
+
+    return vLllmOpenAICompletionFormatModelGenerator(
+        url=url,
+        pat_token=pat_token,
+        prompt_formatter=doc_qa_task_prompt_template,
+        batch_size=1,
+        model_name=model_name,
+        format_prompt_func=vicuna_prompt_format_func,
+        concurrency=100,
+    )
+
+
+def vllm_llama2_model_generator(url, pat_token, model_name):
+    from databricks.labs.doc_qa.model_generators.model_generator import (
+        vLllmOpenAICompletionFormatModelGenerator,
+    )
+    from databricks.labs.doc_qa.variables.doc_qa_template_variables import (
+        llama2_prompt_format_func,
+        doc_qa_task_prompt_template,
+    )
+
+    return vLllmOpenAICompletionFormatModelGenerator(
+        url=url,
+        pat_token=pat_token,
+        prompt_formatter=doc_qa_task_prompt_template,
+        batch_size=1,
+        model_name=model_name,
+        format_prompt_func=llama2_prompt_format_func,
+        concurrency=100,
+    )
+
+
+def generate_and_evaluate(
+    input_df, model_generator, evaluator, temperature=0, max_tokens=200
+):
+    generate_result = model_generator.run_tasks(
+        input_df=input_df, temperature=temperature, max_tokens=max_tokens
+    )
+
+    result_df = generate_result.to_dataframe()
+
+    logger.info(f"Finished generating {len(result_df)} rows, starting evaluation")
+    return evaluator.run_eval(dataset_df=result_df, concurrency=20, catch_error=True)
+
+
+def evaluate_using_vllm_locally(
+    input_df,
+    hf_model_name,
+    prompt_tempate_format_func,
+    temperature=0,
+    max_tokens=200,
+    max_num_batched_tokens=None,
+):
+    from databricks.labs.doc_qa.model_generators.model_generator import (
+        vLllmLocalModelGenerator,
+    )
+    from databricks.labs.doc_qa.variables.doc_qa_template_variables import (
+        doc_qa_task_prompt_template,
+    )
+
+    model_generator = vLllmLocalModelGenerator(
+        hf_model_name=hf_model_name,
+        format_prompt_func=prompt_tempate_format_func,
+        prompt_formatter=doc_qa_task_prompt_template,
+        max_num_batched_tokens=max_num_batched_tokens,
+    )
+
+    evaluator = gpt_4_evaluator()
+    generate_result = model_generator.run_tasks(
+        input_df=input_df, temperature=temperature, max_tokens=max_tokens
+    )
+
+    result_df = generate_result.to_dataframe()
+
+    logger.info(f"Finished generating {len(result_df)} rows, starting evaluation")
+    return evaluator.run_eval(dataset_df=result_df, concurrency=20, catch_error=True)
diff --git a/databricks/labs/doc_qa/evaluators/templated_evaluator.py b/databricks/labs/doc_qa/evaluators/templated_evaluator.py
@@ -7,17 +7,13 @@
 from databricks.labs.doc_qa.llm_providers import openai_provider
 from databricks.labs.doc_qa.llm_providers import anthropic_provider
 import json
-import logging
+from databricks.labs.doc_qa.logging_utils import logger
 from tenacity import retry, stop_after_attempt, retry_if_result, retry_if_exception
 import re
 from enum import Enum
 import json
 
 
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__.split(".")[0])
-
-
 class ParameterType(Enum):
     STRING = 'string'
     NUMBER = 'number'
diff --git a/databricks/labs/doc_qa/llm_providers/anthropic_provider.py b/databricks/labs/doc_qa/llm_providers/anthropic_provider.py
@@ -3,8 +3,9 @@
 import os
 import requests
 from tenacity import retry, stop_after_attempt, stop_after_delay, wait_fixed, retry_if_exception_type, retry_if_exception
-import logging 
+from databricks.labs.doc_qa.logging_utils import logger
 from anthropic import Anthropic, HUMAN_PROMPT, AI_PROMPT
+import logging
 
 anthropic_api_key = os.getenv('ANTHROPIC_API_KEY')
 
@@ -16,7 +17,6 @@ def supress_httpx_logs():
     logger.setLevel(logging.WARNING)
 
 supress_httpx_logs()
-logger = logging.getLogger(__name__)
 
 
 def request_anthropic(prompt, temperature=0.0, model="claude-2", max_tokens_to_sample=300):
diff --git a/databricks/labs/doc_qa/llm_providers/openai_provider.py b/databricks/labs/doc_qa/llm_providers/openai_provider.py
@@ -3,10 +3,7 @@
 import os
 import requests
 from tenacity import retry, stop_after_attempt, stop_after_delay, wait_fixed, retry_if_exception_type, retry_if_exception
-import logging 
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+from databricks.labs.doc_qa.logging_utils import logger
 
 
 openai_token = os.getenv('OPENAI_API_KEY')
diff --git a/databricks/labs/doc_qa/logging_utils.py b/databricks/labs/doc_qa/logging_utils.py
@@ -0,0 +1,10 @@
+import logging
+
+logger = logging.getLogger("doc-qa")
+formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+
+handler = logging.StreamHandler()
+handler.setFormatter(formatter)
+
+logger.addHandler(handler)
+logger.setLevel(logging.INFO)
diff --git a/databricks/labs/doc_qa/model_generators/model_generator.py b/databricks/labs/doc_qa/model_generators/model_generator.py
diff --git a/databricks/labs/doc_qa/variables/doc_qa_template_variables.py b/databricks/labs/doc_qa/variables/doc_qa_template_variables.py
diff --git a/setup.py b/setup.py
