diff --git a/natural_language_processing/embedding_generation/langchain/run.py b/natural_language_processing/embedding_generation/langchain/run.py new file mode 100644 index 00000000..855bfa1c --- /dev/null +++ b/natural_language_processing/embedding_generation/langchain/run.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2024, Ampere Computing LLC +try: + from utils import misc # noqa +except ModuleNotFoundError: + import os + import sys + filename = "set_env_variables.sh" + directory = os.path.realpath(__file__).split("/")[:-1] + for idx in range(1, len(directory) - 1): + subdir = "/".join(directory[:-idx]) + if filename in os.listdir(subdir): + print(f"\nPlease run \033[91m'source {os.path.join(subdir, filename)}'\033[0m first.") + break + else: + print(f"\n\033[91mFAIL: Couldn't find {filename}, are you running this script as part of Ampere Model Library?" + f"\033[0m") + sys.exit(1) + +def run_pytorch_fp32(model_name, num_runs, timeout, filepath, **kwargs): + import os + from langchain_community.document_loaders import TextLoader + from langchain_community.embeddings import HuggingFaceEmbeddings + from langchain.text_splitter import CharacterTextSplitter + from langchain_community.vectorstores import Chroma + from utils.benchmark import run_model + from utils.pytorch import PyTorchRunnerV2, apply_compile + from utils.helpers import DummyDataset + + embedding_model = HuggingFaceEmbeddings(model_name=model_name, show_progress=False) + embedding_model.client.eval() + embedding_model.client.forward = apply_compile(embedding_model.client.forward) + + documents = TextLoader(filepath).load_and_split(CharacterTextSplitter()) + + def single_pass_pytorch(_runner, _): + _runner.run(os.path.getsize(filepath), documents) + + def embeddings_gen(_documents): + return Chroma.from_documents(_documents, embedding_model) + + runner = PyTorchRunnerV2(embeddings_gen, throughput_only=True) + return run_model(single_pass_pytorch, runner, DummyDataset(), 1, num_runs, timeout) + + +if __name__ == "__main__": + from utils.helpers import DefaultArgParser + parser = DefaultArgParser(["pytorch"]) + parser.require_model_name([ + "BAAI/bge-small-en-v1.5", "BAAI/bge-base-en-v1.5", "sentence-transformers/all-MiniLM-L6-v2"]) + parser.add_argument("--filepath", type=str, required=True, help="path to a .txt file") + run_pytorch_fp32(**vars(parser.parse())) diff --git a/utils/helpers.py b/utils/helpers.py index 67fadbf1..0e281a2e 100644 --- a/utils/helpers.py +++ b/utils/helpers.py @@ -77,3 +77,15 @@ def print_accuracy_metrics(self) -> dict: print(f"{3 * indent}{metric[:max_len]}{(max_len - len(metric)) * ' '}{3 * indent}" + "= {:>7.3f}".format(float(accuracy_results[metric]))) return accuracy_results + + +class DummyDataset(Dataset): + def __init__(self): + import sys + self.available_instances = int(sys.maxsize) + + def reset(self) -> bool: + return True + + def _summarize_accuracy(self) -> dict: + return {}