merge main into prestaing

.github/workflows/python-package.yml

@@ -25,7 +25,7 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         python -m pip install flake8 pytest black
-        pip install -e .        
+        bash install.sh
 
     - name: Lint with flake8
       run: |

install.sh

@@ -13,4 +13,4 @@ pip install -r requirements.txt
 pip uninstall uvloop -y
 
 # Reinstalling pydantic and transformers with specific versions that work with our repository and vllm
-pip install pydantic==1.10.7 transformers==4.36.2
+pip install pydantic==1.10.7 transformers==4.36.2 angle_emb==0.3.8 peft==0.9.0

prompting/__init__.py

@@ -16,7 +16,7 @@
 # DEALINGS IN THE SOFTWARE.
 
 # Define the version of the template module.
-__version__ = "1.2.0"
+__version__ = "1.3.0"
 version_split = __version__.split(".")
 __spec_version__ = (
     (10000 * int(version_split[0]))

prompting/conversation.py

@@ -16,7 +16,7 @@
 from transformers import Pipeline
 
 
-def create_task(llm_pipeline: Pipeline, task_name: str) -> Task:
+def create_task(llm_pipeline: Pipeline, task_name: str, create_reference=True) -> Task:
     wiki_based_tasks = ["summarization", "qa"]
     coding_based_tasks = ["debugging"]
     # TODO: Abstract dataset classes into common dynamic interface
@@ -33,20 +33,38 @@ def create_task(llm_pipeline: Pipeline, task_name: str) -> Task:
         dataset = WikiDateDataset()
 
     if task_name == "summarization":
-        task = SummarizationTask(llm_pipeline=llm_pipeline, context=dataset.next())
+        task = SummarizationTask(
+            llm_pipeline=llm_pipeline,
+            context=dataset.next(),
+            create_reference=create_reference,
+        )
 
     elif task_name == "qa":
-        task = QuestionAnsweringTask(llm_pipeline=llm_pipeline, context=dataset.next())
+        task = QuestionAnsweringTask(
+            llm_pipeline=llm_pipeline,
+            context=dataset.next(),
+            create_reference=create_reference,
+        )
 
     elif task_name == "debugging":
-        task = DebuggingTask(llm_pipeline=llm_pipeline, context=dataset.next())
+        task = DebuggingTask(
+            llm_pipeline=llm_pipeline,
+            context=dataset.next(),
+            create_reference=create_reference,
+        )
 
     elif task_name == "math":
-        task = MathTask(llm_pipeline=llm_pipeline, context=dataset.next())
+        task = MathTask(
+            llm_pipeline=llm_pipeline,
+            context=dataset.next(),
+            create_reference=create_reference,
+        )
 
     elif task_name == "date_qa":
         task = DateQuestionAnsweringTask(
-            llm_pipeline=llm_pipeline, context=dataset.next()
+            llm_pipeline=llm_pipeline,
+            context=dataset.next(),
+            create_reference=create_reference,
         )
 
     else:

prompting/forward.py

@@ -18,6 +18,7 @@
 
 import time
 import sys
+import asyncio
 import numpy as np
 import bittensor as bt
 
@@ -29,6 +30,22 @@
 from prompting.rewards import RewardResult
 from prompting.utils.uids import get_random_uids
 from prompting.utils.logging import log_event
+from prompting.utils.misc import async_log
+
+
+@async_log
+async def generate_reference(agent):
+    loop = asyncio.get_running_loop()
+    result = await loop.run_in_executor(
+        None, agent.task.generate_reference, agent.llm_pipeline
+    )
+    return result
+
+
+@async_log
+async def execute_dendrite_call(dendrite_call):
+    responses = await dendrite_call
+    return responses
 
 
 async def handle_response(responses: Dict[int, Awaitable]) -> List[bt.Synapse]:
@@ -137,13 +154,12 @@ async def run_step(
         **response_event.__state_dict__(),
     }
 
-    log_event(self, event)
-
     return event
 
 
 async def forward(self):
     bt.logging.info("🚀 Starting forward loop...")
+    forward_start_time = time.time()
 
     while True:
         bt.logging.info(
@@ -155,7 +171,11 @@ async def forward(self):
         )
         bt.logging.info(f"📋 Creating {task_name} task... ")
         try:
-            task = create_task(llm_pipeline=self.llm_pipeline, task_name=task_name)
+            task = create_task(
+                llm_pipeline=self.llm_pipeline,
+                task_name=task_name,
+                create_reference=False,
+            )
             break
         except Exception as e:
             bt.logging.error(
@@ -180,6 +200,11 @@ async def forward(self):
             timeout=self.config.neuron.timeout,
             exclude=exclude_uids,
         )
+
+        # Adds forward time to event and logs it to wandb
+        event["forward_time"] = time.time() - forward_start_time
+        log_event(self, event)
+
         exclude_uids += event["uids"]
         task.complete = True
 

prompting/llms/vllm_llm.py

@@ -154,14 +154,18 @@ def _make_prompt(self, messages: List[Dict[str, str]]):
 
         for message in messages:
             if message["role"] == "system":
-                composed_prompt += f'<|system|>{message["content"]} '
+                composed_prompt += (
+                    f'<|im_start|>system\n{message["content"]} <|im_end|>'
+                )
             elif message["role"] == "user":
-                composed_prompt += f'<|user|>{message["content"]} '
+                composed_prompt += f'<|im_start|>user\n{message["content"]} <|im_end|>'
             elif message["role"] == "assistant":
-                composed_prompt += f'<|assistant|>{message["content"]} '
+                composed_prompt += (
+                    f'<|im_start|>assistant\n{message["content"]} <|im_end|>'
+                )
 
         # Adds final tag indicating the assistant's turn
-        composed_prompt += "<|assistant|>"
+        composed_prompt += "<|im_start|>assistant\n"
         return composed_prompt
 
     def forward(self, messages: List[Dict[str, str]]):

prompting/tools/datasets/math.py

@@ -57,7 +57,7 @@ def get(
         """
         bt.logging.info(f"Getting math problem {name!r}")
         info = mathgenerator.generate_context(name, **kwargs)
-        if info["reward_type"] != "float":
+        if info["reward_type"] != "float" or info["topic"] == "computer_science":
             return None
 
         math_words = [

prompting/utils/config.py

@@ -269,7 +269,7 @@ def add_validator_args(cls, parser):
         "--neuron.model_id",
         type=str,
         help="The model to use for the validator.",
-        default="HuggingFaceH4/zephyr-7b-beta",
+        default="NousResearch/Nous-Hermes-2-SOLAR-10.7B",
     )
 
     parser.add_argument(

prompting/utils/misc.py

@@ -17,6 +17,9 @@
 # DEALINGS IN THE SOFTWARE.
 
 import time
+import asyncio
+import threading
+import bittensor as bt
 from math import floor
 from typing import Callable, Any
 from functools import lru_cache, update_wrapper
@@ -108,3 +111,24 @@ def ttl_get_block(self) -> int:
     Note: self here is the miner or validator instance
     """
     return self.subtensor.get_current_block()
+
+
+def async_log(func):
+    async def wrapper(*args, **kwargs):
+        start_time = time.time()
+        task_id = id(asyncio.current_task())
+        func_name = func.__name__
+        bt.logging.debug(f"Starting {func_name} on task {task_id} at {start_time}")
+
+        # Execute the wrapped function
+        result = await func(*args, **kwargs)
+
+        end_time = time.time()
+        execution_time = end_time - start_time
+        bt.logging.debug(
+            f"Completed {func_name} on task {task_id} in {execution_time} seconds"
+        )
+
+        return result
+
+    return wrapper

scripts/setup_ubuntu_machine.sh

@@ -23,10 +23,6 @@ git clone https://github.com/opentensor/prompting.git
 # Change to the prompting directory
 cd prompting
 
-# Install Python dependencies
-python3 -m pip install -r requirements.txt
-
-# Install prompting package
-python3 -m pip install -e .
+bash install.sh
 
 echo "Script completed successfully."

tests/test_forward.py

@@ -0,0 +1,59 @@
+import pytest
+import time
+import asyncio
+import sys
+from functools import partial
+import bittensor as bt
+from prompting.forward import run_step
+from neurons.validator import Validator
+from prompting.tasks import Task, QuestionAnsweringTask
+from .fixtures.task import WIKI_CONTEXT
+from prompting.agent import HumanAgent
+from unittest.mock import patch, Mock
+from prompting.protocol import PromptingSynapse
+
+sys.argv = [__file__, "--mock", "--wandb.off", "--neuron.tasks", "qa"]
+mock_neuron = Validator()
+
+task = QuestionAnsweringTask(
+    llm_pipeline=mock_neuron.llm_pipeline, context=WIKI_CONTEXT, create_reference=False
+)
+
+
+def generate_reference(x, delay=1):
+    time.sleep(delay)
+    return "Fake reference"
+
+
+async def mock_dendrite_call(delay=1, **kwargs):
+    time.sleep(delay)
+
+    mock_synapse = PromptingSynapse(roles=["user"], messages=[""])
+    mock_synapse.completion = "Fake response"
+
+    return [mock_synapse]
+
+
+@pytest.mark.parametrize(
+    "generate_reference_time, dendrite_time, expected_forward_time",
+    [(0.5, 0.5, 0.5), (0.5, 0.4, 0.5), (0.4, 0.5, 0.5)],
+)
+def test_generate_reference_parallel_to_dendrite(
+    generate_reference_time, dendrite_time, expected_forward_time
+):
+    task.generate_reference = partial(generate_reference, delay=generate_reference_time)
+    mock_agent = HumanAgent(task, mock_neuron.llm_pipeline)
+
+    mock_neuron.dendrite = partial(mock_dendrite_call, delay=dendrite_time)
+
+    event = asyncio.run(run_step(mock_neuron, mock_agent, k=4, timeout=0.1))
+
+    step_time = event["step_time"]
+    reward_pipeline_time = sum(
+        event[key] for key in event if key.endswith("batch_time")
+    )
+    network_and_reference_gen_time = step_time - reward_pipeline_time
+
+    assert network_and_reference_gen_time == pytest.approx(
+        expected_forward_time, abs=0.1
+    )

tests/test_mock.py

@@ -76,6 +76,7 @@ async def run():
             timeout=timeout,
         )
 
+    eps = 0.2
     responses = asyncio.run(run())
     for synapse in responses:
         assert (
@@ -89,9 +90,9 @@ async def run():
 
         # check that the dendrite take between min_time and max_time
         assert min_time <= dendrite.process_time
-        assert dendrite.process_time <= max_time + 0.1
+        assert dendrite.process_time <= max_time + eps
         # check that responses which take longer than timeout have 408 status code
-        if dendrite.process_time >= timeout + 0.1:
+        if dendrite.process_time >= timeout + eps:
             assert dendrite.status_code == 408
             assert dendrite.status_message == "Timeout"
             assert synapse.completion == ""