Formatting + sorting imports

.github/workflows/style-check.yml

@@ -0,0 +1,35 @@
+# Based on https://github.com/psf/black/actions/runs/17913292/workflow
+
+name: Style check
+on:
+  push:
+    paths:
+      - '**.py'
+  pull_request:
+    paths:
+      - '**.py'
+jobs:
+  black:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+      - name: Install Black
+        run: pip install black
+      - name: Run black check
+        run: black --check .
+  isort:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+      - name: Install isort
+        run: pip install isort
+      - name: Run isort check
+        run: isort --check-only .

example_configure_survey.py

@@ -1,6 +1,6 @@
-from lm_survey.survey.survey import Survey
 import os
 
+from lm_survey.survey.survey import Survey
 
 if __name__ == "__main__":
     survey_directory = os.path.join("data", "ATP", "American_Trends_Panel_W92")

lm_survey/constants.py

@@ -1,6 +1,5 @@
 import typing
 
-
 MULTIPLE_CHOICE_LIST: typing.List[str] = [
     "A",
     "B",

lm_survey/prompt_templates.py

@@ -2,7 +2,6 @@
 
 from lm_survey.constants import MULTIPLE_CHOICE_LIST
 
-
 INDEPENDENT_VARIABLE_SUMMARY_TEMPLATE = """Context: {context_summary}
 
 {dependent_variable_prompt}"""

lm_survey/samplers/__init__.py

@@ -1,5 +1,5 @@
-from lm_survey.samplers.base_sampler import BaseSampler
+from lm_survey.samplers.async_openai_sampler import AsyncOpenAiSampler
 from lm_survey.samplers.auto_sampler import AutoSampler
+from lm_survey.samplers.base_sampler import BaseSampler
 from lm_survey.samplers.hf_sampler import HfSampler
 from lm_survey.samplers.openai_sampler import OpenAiSampler
-from lm_survey.samplers.async_openai_sampler import AsyncOpenAiSampler

lm_survey/samplers/async_openai_sampler.py

@@ -1,10 +1,9 @@
 import sys
 
 import openai
-from openai.error import RateLimitError
-
 import torch
 from aiolimiter import AsyncLimiter
+from openai.error import RateLimitError
 
 from lm_survey.samplers.base_sampler import BaseSampler, MaybeAwaitable
 

lm_survey/samplers/auto_sampler.py

@@ -1,7 +1,7 @@
-from lm_survey.samplers.hf_sampler import HfSampler
-from lm_survey.samplers.openai_sampler import OpenAiSampler
 from lm_survey.samplers.async_openai_sampler import AsyncOpenAiSampler
 from lm_survey.samplers.base_sampler import BaseSampler
+from lm_survey.samplers.hf_sampler import HfSampler
+from lm_survey.samplers.openai_sampler import OpenAiSampler
 
 
 class AutoSampler(BaseSampler):

lm_survey/samplers/base_sampler.py

@@ -1,5 +1,5 @@
-from abc import ABCMeta, abstractmethod
 import typing
+from abc import ABCMeta, abstractmethod
 
 T = typing.TypeVar("T")
 MaybeAwaitable = typing.Union[T, typing.Awaitable[T]]

lm_survey/samplers/hf_sampler.py

@@ -1,8 +1,9 @@
 import typing
-from lm_survey.samplers.base_sampler import BaseSampler
 
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from lm_survey.samplers.base_sampler import BaseSampler
 
 
 class HfSampler(BaseSampler):

lm_survey/samplers/openai_sampler.py

@@ -1,10 +1,10 @@
 import typing
 
+import openai
 import tiktoken
 import torch
 
 from lm_survey.samplers.base_sampler import BaseSampler
-import openai
 
 OPENAI_TOKEN_COSTS = {
     # cents per 1000 tokens

lm_survey/survey/__init__.py

@@ -1,8 +1,8 @@
 from lm_survey.survey.dependent_variable_sample import (
-    DependentVariableSample,
     Completion,
+    DependentVariableSample,
 )
-from lm_survey.survey.survey import Survey
-from lm_survey.survey.variable import Variable
 from lm_survey.survey.question import Question, ValidOption
 from lm_survey.survey.results import SurveyResults
+from lm_survey.survey.survey import Survey
+from lm_survey.survey.variable import Variable

lm_survey/survey/dependent_variable_sample.py

@@ -1,4 +1,5 @@
 import typing
+
 from lm_survey.survey.question import Question
 
 

lm_survey/survey/question.py

@@ -1,5 +1,7 @@
 import typing
+
 import pandas as pd
+
 from lm_survey.constants import MULTIPLE_CHOICE_LIST
 from lm_survey.prompt_templates import (
     DEPENDENT_VARIABLE_TEMPLATE,

lm_survey/survey/results.py

@@ -1,8 +1,10 @@
 import json
 import os
 import typing
+
 import pandas as pd
 import pandas.core.groupby.generic
+
 from lm_survey.survey.dependent_variable_sample import DependentVariableSample
 
 

lm_survey/survey/survey.py

@@ -1,3 +1,6 @@
+import argparse
+import functools
+import json
 import os
 import typing
 from pathlib import Path
@@ -6,16 +9,13 @@
 import pandas as pd
 from sklearn.metrics import normalized_mutual_info_score
 
+from lm_survey.prompt_templates import INDEPENDENT_VARIABLE_SUMMARY_TEMPLATE
 from lm_survey.survey.dependent_variable_sample import (
     Completion,
     DependentVariableSample,
 )
 from lm_survey.survey.question import Question, ValidOption
 from lm_survey.survey.variable import Variable
-from lm_survey.prompt_templates import INDEPENDENT_VARIABLE_SUMMARY_TEMPLATE
-import json
-import functools
-import argparse
 
 
 class Survey:

lm_survey/survey/variable.py

@@ -1,5 +1,7 @@
 import typing
+
 import pandas as pd
+
 from lm_survey.survey.question import Question
 
 

pyproject.toml

@@ -0,0 +1,2 @@
+[tool.isort]
+profile = "black"