diff --git a/.github/workflows/build_push_image.yml b/.github/workflows/build_push_image.yml index a10eb8c..512f0c0 100644 --- a/.github/workflows/build_push_image.yml +++ b/.github/workflows/build_push_image.yml @@ -48,6 +48,6 @@ jobs: cd src/compute_horde_prompt_gen ls - docker build -t "${{ IMAGE_NAME }}" . + docker build -t $IMAGE_NAME . - # docker run -v ./output/:/app/output/ "${{ IMAGE_NAME }}" --number_of_batches 5 --number_of_prompts_per_batch 30 --uuids uuid1,uuid2,uuid3,uuid4,uuid5 + # docker run -v ./output/:/app/output/ $IMAGE_NAME --number_of_batches 5 --number_of_prompts_per_batch 30 --uuids uuid1,uuid2,uuid3,uuid4,uuid5 diff --git a/.github/workflows/smoke_test.yml b/.github/workflows/smoke_test.yml index d9efd2b..8428fd5 100644 --- a/.github/workflows/smoke_test.yml +++ b/.github/workflows/smoke_test.yml @@ -23,11 +23,11 @@ jobs: with: python-version: ${{ env.PYTHON_DEFAULT_VERSION }} - - name: Install dependencies - run: | - # python -m pip install --upgrade 'pdm>=2.12,<3' - # pdm install - python -m pip install transformers torch + # - name: Install dependencies + # run: | + # # python -m pip install --upgrade 'pdm>=2.12,<3' + # # pdm install + # python -m pip install transformers torch - name: Run Test run: | diff --git a/src/compute_horde_prompt_gen/model.py b/src/compute_horde_prompt_gen/model.py index 4998646..8245ab9 100644 --- a/src/compute_horde_prompt_gen/model.py +++ b/src/compute_horde_prompt_gen/model.py @@ -1,9 +1,5 @@ -import torch +import numpy as np import logging -from transformers import ( - AutoTokenizer, - AutoModelForCausalLM, -) from prompt import PROMPT_ENDING @@ -15,7 +11,7 @@ def __init__(self): pass def generate(self, prompts: list[str], num_return_sequences: int, **_kwargs): - return torch.rand(len(prompts) * num_return_sequences) + return np.random.rand(len(prompts) * num_return_sequences) def decode(self, _output): return f"COPY PASTE INPUT PROMPT {PROMPT_ENDING} Here is the list of prompts:\nHow are you?\nDescribe something\nCount to ten\n" @@ -23,6 +19,12 @@ def decode(self, _output): class GenerativeModel: def __init__(self, model_path: str, quantize: bool = False): + import torch + from transformers import ( + AutoTokenizer, + AutoModelForCausalLM, + ) + quantization_config = None if quantize: from transformers import BitsAndBytesConfig