Stylized responses (aka responses with custom system prompt).

airoboros/instructors/agent.py

@@ -1,9 +1,9 @@
 from airoboros.instructors.inline_qa import generate as generate_inline
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for agent/router training data."""
     async for item in generate_inline(
-        instructor, "agent", start_key="PROMPT", filter_response=False
+        instructor, "agent", start_key="PROMPT", filter_response=False, **kwargs
     ):
         yield item

airoboros/instructors/card.py

@@ -1,7 +1,7 @@
 from airoboros.instructors.inline_qa import generate as generate_inline
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for model/scenario card training data."""
-    async for item in generate_inline(instructor, "card", start_key="PROMPT"):
+    async for item in generate_inline(instructor, "card", start_key="PROMPT", **kwargs):
         yield item

airoboros/instructors/chat.py

@@ -171,6 +171,7 @@ async def generate_cards(instructor):
             logger.success(f"Generated chat character card {filename}")
             if len(cards) >= card_count:
                 break
+    instructor.instructor_counts["chat_card"] = len(cards)
     return cards
 
 
@@ -427,7 +428,7 @@ async def generate_chat(instructor, cards, topic, **api_params):
     return training
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for chat training data."""
     config = instructor.instructors.get("chat", {})
     if not config:

airoboros/instructors/coding.py

@@ -4,7 +4,7 @@
 import random
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for coding training data."""
     config = instructor.instructors.get("coding")
     if not config:

airoboros/instructors/contextual.py

@@ -139,7 +139,7 @@ def generate_prompt(instructor, config, template, topic_iter):
     )
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for contextual training data."""
     config = instructor.instructors.get("contextual")
     if not config:

airoboros/instructors/cot.py

@@ -1,7 +1,9 @@
 from airoboros.instructors.simple_task import generate as generate_simple
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for chain-of-thought training data."""
-    async for item in generate_simple(instructor, "cot", filter_response=False):
+    async for item in generate_simple(
+        instructor, "cot", filter_response=False, **kwargs
+    ):
         yield item

airoboros/instructors/counterfactual_contextual.py

@@ -19,7 +19,7 @@
 ]
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for contextual training data."""
     config = instructor.instructors.get("counterfactual_contextual")
     if not config:

airoboros/instructors/detailed_writing.py

@@ -24,7 +24,7 @@
 )
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for detailed writing training data."""
     config = instructor.instructors.get("detailed_writing", {})
     if not config:
@@ -60,8 +60,7 @@ async def generate(instructor):
 
     # Load the topics.
     topics = instructor.get_instructor_topics(config)
-    random.shuffle(topics)
-    topic_index = 0
+    topic_index = random.randint(0, len(topics) - 1)
 
     # API params, overriding defaults with this instructor's config.
     api_params = {**instructor.api_params, **config.get("api_params", {})}

airoboros/instructors/experience.py

@@ -3,7 +3,7 @@
 import re
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for experiences."""
     config = instructor.instructors.get("experience")
     if not config:
@@ -43,6 +43,7 @@ async def generate(instructor):
         futures.append(
             instructor.generate_response(
                 prompt.format(language=language, flesch=flesch),
+                messages=kwargs.get("messages", []),
                 filter_response=False,
                 **api_params
             )

airoboros/instructors/general.py

@@ -1,10 +1,11 @@
 import asyncio
 import json
 import os
+import random
 import re
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for generic/general training data."""
     config = instructor.instructors.get("general")
     if not config:
@@ -25,9 +26,7 @@ async def generate(instructor):
 
     # Load the topics.
     topics = instructor.get_instructor_topics(config)
-    topic_index = instructor.instructor_counts.get("general", 0)
-    if topic_index >= len(topics):
-        topic_index = topic_index % len(topics)
+    topic_index = random.randint(0, len(topics) - 1)
 
     # API params, overriding defaults with this instructor's config.
     api_params = {**instructor.api_params, **config.get("api_params", {})}
@@ -85,7 +84,11 @@ async def generate(instructor):
             ):
                 continue
             instructions.append(instruction)
-            futures.append(instructor.generate_response(instruction, **api_params))
+            futures.append(
+                instructor.generate_response(
+                    instruction, messages=kwargs.get("messages", []), **api_params
+                )
+            )
         if not futures:
             continue
         responses = await asyncio.gather(*futures)

airoboros/instructors/inline_qa.py

@@ -1,5 +1,6 @@
 import json
 import os
+import random
 import re
 
 
@@ -10,6 +11,7 @@ async def generate(
     end_key="ANSWER",
     filter_response=True,
     template_kwargs={},
+    **kwargs,
 ):
     """Generator for generic inline question answer pair training data."""
     config = instructor.instructors.get(category)
@@ -31,9 +33,7 @@ async def generate(
 
     # Topics included in template?
     topics = instructor.get_instructor_topics(config)
-    topic_index = instructor.instructor_counts.get(category, 0)
-    if topic_index >= len(topics):
-        topic_index = topic_index % len(topics)
+    topic_index = random.randint(0, len(topics) - 1)
 
     # API params, overriding defaults with this instructor's config.
     api_params = {**instructor.api_params, **config.get("api_params", {})}
@@ -83,7 +83,10 @@ async def generate(
 
         prompt = template.format(**prompt_args)
         response = await instructor.generate_response(
-            prompt, filter_response=filter_response, **api_params
+            prompt,
+            messages=kwargs.get("messages", []),
+            filter_response=filter_response,
+            **api_params,
         )
         if not response:
             continue

airoboros/instructors/joke.py

@@ -1,7 +1,7 @@
 from airoboros.instructors.inline_qa import generate as generate_inline
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for joke training data."""
-    async for item in generate_inline(instructor, "joke"):
+    async for item in generate_inline(instructor, "joke", **kwargs):
         yield item

airoboros/instructors/multiple_choice.py

@@ -13,7 +13,7 @@
 )
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for multiple choice training data."""
     config = instructor.instructors.get("multiple_choice")
     if not config:

airoboros/instructors/orca.py

@@ -1,7 +1,7 @@
 from airoboros.instructors.inline_qa import generate as generate_inline
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for orca training data."""
-    async for item in generate_inline(instructor, "orca"):
+    async for item in generate_inline(instructor, "orca", **kwargs):
         yield item

airoboros/instructors/plan.py

@@ -1,7 +1,9 @@
 from airoboros.instructors.simple_task import generate as generate_simple
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for rewoo style planning."""
-    async for item in generate_simple(instructor, "plan", filter_response=False):
+    async for item in generate_simple(
+        instructor, "plan", filter_response=False, **kwargs
+    ):
         yield item

airoboros/instructors/riddle.py

@@ -1,7 +1,7 @@
 from airoboros.instructors.inline_qa import generate as generate_inline
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for riddle training data."""
-    async for item in generate_inline(instructor, "riddle"):
+    async for item in generate_inline(instructor, "riddle", **kwargs):
         yield item

airoboros/instructors/roleplay.py

@@ -1,7 +1,7 @@
 from airoboros.instructors.simple_task import generate as generate_simple_task
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for roleplay training data."""
-    async for item in generate_simple_task(instructor, "roleplay"):
+    async for item in generate_simple_task(instructor, "roleplay", **kwargs):
         yield item

airoboros/instructors/simple_task.py

@@ -1,5 +1,7 @@
 import asyncio
+import json
 import os
+import random
 import re
 
 
@@ -9,6 +11,7 @@ async def generate(
     filter_response=True,
     only_instructions=False,
     template_kwargs={},
+    **kwargs,
 ):
     """Generator for simple instruction response tasks (e.g. roleplay, wordgames)."""
     config = instructor.instructors.get(category)
@@ -46,6 +49,10 @@ async def generate(
         min_score = instructor.min_docsearch_score
     min_score = float(min_score)
 
+    # Load the topics.
+    topics = instructor.get_instructor_topics(config)
+    topic_index = random.randint(0, len(topics) - 1)
+
     # Generate the instruction/response pairs until we reach the target count.
     batch_size = config.get("batch_size")
     if batch_size is None:
@@ -63,6 +70,21 @@ async def generate(
             format_args[key] = val(instructor)
         if "{topic_avoidance}" in template:
             format_args["topic_avoidance"] = instructor.topic_avoidance
+        if "{topics}" in template:
+            # Inject the topics to use for this batch.
+            current_topics = []
+            for _ in range(batch_size):
+                current_topics.append(topics[topic_index])
+                topic_index += 1
+                if topic_index >= len(topics):
+                    topic_index = 0
+            topics_str = "\n".join(
+                [
+                    f" * TSK {idx + 1} must be related to topic: {json.dumps(topic)}"
+                    for idx, topic in enumerate(current_topics)
+                ]
+            )
+            format_args["topics"] = topics_str
 
         # Get a batch of instructions.
         prompt = template.format(**format_args)
@@ -93,7 +115,10 @@ async def generate(
                     )
                 futures.append(
                     instructor.generate_response(
-                        full_prompt, filter_response=filter_response, **api_params
+                        full_prompt,
+                        messages=kwargs.get("messages", []),
+                        filter_response=filter_response,
+                        **api_params,
                     )
                 )
         if not futures:

airoboros/instructors/song.py

@@ -1,7 +1,7 @@
 from airoboros.instructors.simple_task import generate as generate_simple_task
 
 
-async def generate(instructor):
+async def generate(instructor, **kwargs):
     """Generator for song training data."""
-    async for item in generate_simple_task(instructor, "song"):
+    async for item in generate_simple_task(instructor, "song", **kwargs):
         yield item