Merge remote-tracking branch 'upstream/main'

Signed-off-by: Dushyant Behl <dushyantbehl@users.noreply.github.com>

Author: dushyantbehl

.github/dependabot.yml

@@ -0,0 +1,6 @@
+version: 2
+updates:
+  - package-ecosystem: "pip"
+    directory: "/"
+    schedule:
+      interval: "daily"

.github/workflows/format.yml

@@ -0,0 +1,38 @@
+# Copyright The Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: Format
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python 3.9
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.9
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install -r setup_requirements.txt
+      - name: Check Formatting
+        run: tox -e fmt
+

.gitignore

@@ -0,0 +1,28 @@
+*.egg-info
+*.pyc
+__pycache__
+.coverage
+.coverage.*
+durations/*
+coverage*.xml
+dist
+htmlcov
+build
+test
+
+# IDEs
+.vscode/
+.idea/
+
+# Env files
+.env
+
+# Virtual Env
+venv/
+.venv/
+
+# Mac personalization files
+*.DS_Store
+
+# Tox envs
+.tox

.isort.cfg

@@ -0,0 +1,10 @@
+[settings]
+profile=black
+from_first=true
+import_heading_future=Future
+import_heading_stdlib=Standard
+import_heading_thirdparty=Third Party
+import_heading_firstparty=First Party
+import_heading_localfolder=Local
+known_firstparty=
+known_localfolder=tuning

.pre-commit-config.yaml

@@ -0,0 +1,11 @@
+repos:
+    - repo: https://github.com/psf/black
+      rev: 22.3.0
+      hooks:
+          - id: black
+            exclude: imports
+    - repo: https://github.com/PyCQA/isort
+      rev: 5.11.5
+      hooks:
+          - id: isort
+            exclude: imports

requirements.txt

@@ -8,7 +8,7 @@ tokenizers>=0.13.3
 tqdm
 trl
 ninja
-peft
+peft>=0.8.0
 datasets>=2.15.0
 flash-attn
 fire

scripts/fmt.sh

@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+
+pre-commit run --all-files
+RETURN_CODE=$?
+
+function echoWarning() {
+  LIGHT_YELLOW='\033[1;33m'
+  NC='\033[0m' # No Color
+  echo -e "${LIGHT_YELLOW}${1}${NC}"
+}
+
+if [ "$RETURN_CODE" -ne 0 ]; then
+  if [ "${CI}" != "true" ]; then
+    echoWarning "☝️ This appears to have failed, but actually your files have been formatted."
+    echoWarning "Make a new commit with these changes before making a pull request."
+  else
+    echoWarning "This test failed because your code isn't formatted correctly."
+    echoWarning 'Locally, run `make run fmt`, it will appear to fail, but change files.'
+    echoWarning "Add the changed files to your commit and this stage will pass."
+  fi
+
+  exit $RETURN_CODE
+fi

scripts/run_inference.py

@@ -8,13 +8,16 @@
 
 If these things change in the future, we should consider breaking it up.
 """
+# Standard
 import argparse
 import json
 import os
+
+# Third Party
 from peft import AutoPeftModelForCausalLM
-import torch
 from tqdm import tqdm
 from transformers import AutoTokenizer
+import torch
 
 
 ### Utilities
@@ -30,10 +33,13 @@ class AdapterConfigPatcher:
             # When loaded in this block, the config's base_model_name_or_path is "foo"
             peft_model = AutoPeftModelForCausalLM.from_pretrained(checkpoint_path)
     """
+
     def __init__(self, checkpoint_path: str, overrides: dict):
         self.checkpoint_path = checkpoint_path
         self.overrides = overrides
-        self.config_path = AdapterConfigPatcher._locate_adapter_config(self.checkpoint_path)
+        self.config_path = AdapterConfigPatcher._locate_adapter_config(
+            self.checkpoint_path
+        )
         # Values that we will patch later on
         self.patched_values = {}
 
@@ -58,7 +64,7 @@ def _locate_adapter_config(checkpoint_path: str) -> str:
     def _apply_config_changes(self, overrides: dict) -> dict:
         """Applies a patch to a config with some override dict, returning the values
         that we patched over so that they may be restored later.
-        
+
         Args:
             overrides: dict
                 Overrides to write into the adapter_config.json. Currently, we
@@ -99,7 +105,9 @@ def _get_old_config_values(adapter_config: dict, overrides: dict) -> dict:
         # For now, we only expect to patch the base model; this may change in the future,
         # but ensure that anything we are patching is defined in the original config
         if not set(overrides.keys()).issubset(set(adapter_config.keys())):
-            raise KeyError("Adapter config overrides must be set in the config being patched")
+            raise KeyError(
+                "Adapter config overrides must be set in the config being patched"
+            )
         return {key: adapter_config[key] for key in overrides}
 
     def __enter__(self):
@@ -119,7 +127,9 @@ def __init__(self, model, tokenizer, device):
         self.device = device
 
     @classmethod
-    def load(cls, checkpoint_path: str, base_model_name_or_path: str=None) -> "TunedCausalLM":
+    def load(
+        cls, checkpoint_path: str, base_model_name_or_path: str = None
+    ) -> "TunedCausalLM":
         """Loads an instance of this model.
 
         Args:
@@ -138,7 +148,11 @@ def load(cls, checkpoint_path: str, base_model_name_or_path: str=None) -> "Tuned
             TunedCausalLM
                 An instance of this class on which we can run inference.
         """
-        overrides = {"base_model_name_or_path": base_model_name_or_path} if base_model_name_or_path is not None else {}
+        overrides = (
+            {"base_model_name_or_path": base_model_name_or_path}
+            if base_model_name_or_path is not None
+            else {}
+        )
         tokenizer = AutoTokenizer.from_pretrained(checkpoint_path)
         # Apply the configs to the adapter config of this model; if no overrides
         # are provided, then the context manager doesn't have any effect.
@@ -153,7 +167,6 @@ def load(cls, checkpoint_path: str, base_model_name_or_path: str=None) -> "Tuned
         peft_model.to(device)
         return cls(peft_model, tokenizer, device)
 
-
     def run(self, text: str, *, max_new_tokens: int) -> str:
         """Runs inference on an instance of this model.
 
@@ -165,13 +178,17 @@ def run(self, text: str, *, max_new_tokens: int) -> str:
 
         Returns:
             str
-                Text generation result.          
+                Text generation result.
         """
         tok_res = self.tokenizer(text, return_tensors="pt")
         input_ids = tok_res.input_ids.to(self.device)
 
-        peft_outputs = self.peft_model.generate(input_ids=input_ids, max_new_tokens=max_new_tokens)
-        decoded_result = self.tokenizer.batch_decode(peft_outputs, skip_special_tokens=False)[0]
+        peft_outputs = self.peft_model.generate(
+            input_ids=input_ids, max_new_tokens=max_new_tokens
+        )
+        decoded_result = self.tokenizer.batch_decode(
+            peft_outputs, skip_special_tokens=False
+        )[0]
         return decoded_result
 
 
@@ -180,7 +197,9 @@ def main():
     parser = argparse.ArgumentParser(
         description="Loads a tuned model and runs an inference call(s) through it"
     )
-    parser.add_argument("--model", help="Path to tuned model to be loaded", required=True)
+    parser.add_argument(
+        "--model", help="Path to tuned model to be loaded", required=True
+    )
     parser.add_argument(
         "--out_file",
         help="JSON file to write results to",
@@ -189,7 +208,7 @@ def main():
     parser.add_argument(
         "--base_model_name_or_path",
         help="Override for base model to be used [default: value in model adapter_config.json]",
-        default=None
+        default=None,
     )
     parser.add_argument(
         "--max_new_tokens",
@@ -199,7 +218,10 @@ def main():
     )
     group = parser.add_mutually_exclusive_group(required=True)
     group.add_argument("--text", help="Text to run inference on")
-    group.add_argument("--text_file", help="File to be processed where each line is a text to run inference on")
+    group.add_argument(
+        "--text_file",
+        help="File to be processed where each line is a text to run inference on",
+    )
     args = parser.parse_args()
     # If we passed a file, check if it exists before doing anything else
     if args.text_file and not os.path.isfile(args.text_file):
@@ -220,7 +242,10 @@ def main():
 
     # TODO: we should add batch inference support
     results = [
-        {"input": text, "output": loaded_model.run(text, max_new_tokens=args.max_new_tokens)}
+        {
+            "input": text,
+            "output": loaded_model.run(text, max_new_tokens=args.max_new_tokens),
+        }
         for text in tqdm(texts)
     ]
 
@@ -230,5 +255,6 @@ def main():
 
     print(f"Exported results to: {args.out_file}")
 
+
 if __name__ == "__main__":
     main()

setup.py

@@ -1,7 +1,4 @@
+# Third Party
 from setuptools import find_packages, setup
 
-setup(
-    name="tuning",
-    version="0.0.1",
-    packages=find_packages()
-)
+setup(name="tuning", version="0.0.1", packages=find_packages())

setup_requirements.txt

@@ -0,0 +1,3 @@
+pre-commit>=3.0.4,<4.0
+pydeps>=1.12.12,<2
+tox>=4.4.2,<5