Camera ready

configs/base.yaml

@@ -73,6 +73,7 @@ inference:
     top_p: 0.1
     temperature: 0.6
     model: ${model.checkpoint}
+    num_display_generations: 5
 
 evaluation:
     additional_metrics:

configs/data/aeslc.yaml

@@ -12,3 +12,4 @@ is_in_conversational_format: false
 system_prompt: "Write a subject text for the email.\n\n"
 assistant_response_start: ${model.assistant_response_start}
 use_test_benchmark: false
+task: 'summarization'

configs/data/example.yaml

@@ -5,7 +5,7 @@ few_shot:
   count: 0
 unlabeled_data_split_name: str
 test_split_name: str
-input_column_name: str
+input_column_name: str or dict[str, str]
 input_max_length: int
 output_column_name: str or list[str]
 output_max_length: int
@@ -15,3 +15,4 @@ system_prompt: str
 assistant_response_start: str
 is_in_conversational_format: bool
 use_test_benchmark: bool
+task: 'str'

configs/data/gigaword.yaml

@@ -10,3 +10,4 @@ output_max_length: 21
 fetch_kwargs: {}
 system_prompt: "Generate a headline for the article in lowercase."
 use_test_benchmark: false
+task: 'summarization'

configs/data/gsm8k.yaml

@@ -0,0 +1,15 @@
+dataset: ['openai/gsm8k', 'main']
+input_column_name: 'question'
+output_column_name: 'answer'
+unlabeled_data_split_name: train
+test_split_name: test
+train_subset_size: null
+test_subset_size: null
+input_max_length: 2048
+output_max_length: 1024
+fetch_kwargs: {}
+is_in_conversational_format: false
+system_prompt: "Please generate a solution to the math problem below. In the end, add a line break and print the answer with \"#### <answer>\" replacing the <answer> tag with the answer."
+assistant_response_start: ${model.assistant_response_start}
+use_test_benchmark: false
+task: 'math'

configs/data/math.yaml

@@ -0,0 +1,15 @@
+dataset: 'Aktsvigun/preprocessed_for_al_OpenR1-Math-220k'
+input_column_name: 'problem'
+output_column_name: {'train': 'short_solution', 'test': 'answer'}
+unlabeled_data_split_name: train
+test_split_name: test
+train_subset_size: null
+test_subset_size: 1_000
+input_max_length: 4096
+output_max_length: 4096
+fetch_kwargs: {}
+is_in_conversational_format: false
+system_prompt: "Please solve the math problem. In the end, add two line breaks and print the answer with \"Answer: <answer>\"."
+assistant_response_start: ${model.assistant_response_start}
+use_test_benchmark: false
+task: 'math'

configs/data/race.yaml

@@ -0,0 +1,27 @@
+dataset: ['ehovy/race', 'all']
+input_column_name: {"article": "article", "question": "question", "options": "options"}
+output_column_name: 'answer'
+unlabeled_data_split_name: train
+test_split_name: test
+train_subset_size: null
+test_subset_size: 1_000
+input_max_length: 1024
+output_max_length: 1
+fetch_kwargs: {}
+is_in_conversational_format: false
+system_prompt: Please answer the question based on the article. Answer only with the letter corresponding to the correct answer. Only output the letter of the correct answer and nothing else.
+user_prompt_template: |
+    Article:
+    {article}
+
+    Question:
+    {question}
+
+    Options:
+    {options}
+
+    Correct answer (only the letter):
+
+assistant_response_start: ${model.assistant_response_start}
+use_test_benchmark: false
+task: 'multi-choice-qa'

configs/data/test.yaml

@@ -14,3 +14,4 @@ train_subset_size: 100
 is_in_conversational_format: false
 assistant_response_start: ${model.assistant_response_start}
 use_test_benchmark: false
+task: 'summarization'

configs/data/trivia_qa.yaml

@@ -12,3 +12,4 @@ is_in_conversational_format: false
 system_prompt: "Please answer the question very laconically:\n\n"
 assistant_response_start: ${model.assistant_response_start}
 use_test_benchmark: false
+task: 'open-qa'

configs/data/user_data.yaml

@@ -11,3 +11,4 @@ system_prompt: "Write a summary for the article in lowercase."
 test_subset_size: 100
 train_subset_size: 100
 use_test_benchmark: false
+task: 'summarization'

configs/data/xsum.yaml

@@ -12,3 +12,4 @@ fetch_kwargs:
 is_in_conversational_format: false
 system_prompt: "Summarize the article in one sentence."
 use_test_benchmark: false
+task: 'summarization'

configs/labeller/api_llm.yaml

@@ -1,6 +1,7 @@
 type: api_llm
 provider: openai
 api_key: 'api_key.key'
+base_url: null
 mode: sync  # or batched
 parameters:
   model: gpt-4o

configs/test.yaml

@@ -62,6 +62,7 @@ inference:
     temperature: 0.6
     top_p: 0.1
     model: ${model.checkpoint}
+    num_display_generations: 0
 
 evaluation:
     additional_metrics: []

src/atgen/labellers/api_labellers/anthropic_labeller.py

@@ -31,9 +31,9 @@ def __call__(self, dataset: Dataset) -> Dataset:
 
         annotations = []
         price = 0
-        for text in tqdm(data):
+        for messages in tqdm(data):
             request_kwargs = dict(base_request_kwargs)
-            request_kwargs["messages"] = [{"role": "user", "content": text}]
+            request_kwargs["messages"] = messages
             output = self.client.messages.create(**request_kwargs)
             price += self._calculate_price(output)
             annotations.append(output.content[0].text.strip())

src/atgen/labellers/api_labellers/openai_labeller.py

@@ -22,38 +22,6 @@
 MAX_NUM_TRIES = 3
 UPDATE_TIME_IN_SECONDS = 10  # update time when checking for the completion
 
-
-# SYSTEM_PROMPT = """
-# Act as an experienced ...
-# [INSERT]
-# Examples for few-shot learning are below.
-# """.strip()
-
-# EXAMPLE_1_INPUT = """
-# [INSERT]
-# """.lstrip()
-#
-# EXAMPLE_2_INPUT = """
-# [INSERT]
-# """.lstrip()
-#
-# EXAMPLE_1_OUTPUT = """
-# [INSERT]
-# """.strip()
-#
-# EXAMPLE_2_OUTPUT = """
-# [INSERT]
-# """.strip()
-
-# Can remove / add examples. If you remove them completely, do not forget to remove the last line from the system prompt
-messages_template = [
-    # {"role": "system", "content": SYSTEM_PROMPT},
-    # {"role": "system", "name": "example_user", "content": EXAMPLE_1_INPUT},
-    # {"role": "system", "name": "example_assistant", "content": EXAMPLE_1_OUTPUT},
-    {"role": "user", "content": ""},
-]
-
-
 class OpenAILabeller(BaseLabeler):
     def __init__(
         self,
@@ -65,7 +33,10 @@ def __init__(
         super().__init__(output_column_name, budget)
         self.config = config
         # Create the OpenAI client
-        kwargs = {} if base_url is None else {"base_url": base_url}
+        if base_url := (base_url or config.base_url):
+            kwargs = {"base_url": base_url}
+        else:
+            kwargs = {}
         self.client = OpenAI(api_key=self.config.api_key, **kwargs)
         self.mode = config.get("mode")
 
@@ -75,16 +46,12 @@ def _sync_call(self, dataset: Dataset) -> Dataset:
         annotations = []
         total_price = 0
 
-        for text in tqdm(data, desc="Processing with OpenAI API"):
-            # Prepare messages for this input
-            text_messages = deepcopy(messages_template)
-            text_messages[-1]["content"] = text
-
+        for messages in tqdm(data, desc="Processing with OpenAI API"):
             # Make API call with retries
             for attempt in range(MAX_NUM_TRIES):
                 try:
                     response = self.client.chat.completions.create(
-                        messages=text_messages, **self.config.parameters
+                        messages=messages, **self.config.parameters
                     )
                     break
                 except Exception as e:
@@ -135,6 +102,10 @@ def _sync_call(self, dataset: Dataset) -> Dataset:
             dataset = dataset.remove_columns(self.output_column_name)
 
         dataset = dataset.add_column(self.output_column_name, annotations)
+        # TODO: make a parameter
+        for i, annotation in enumerate(annotations[:2]):
+            print(f"Annotation of the {i+1}th instance: {annotation}")
+            print("-"*100)
         return dataset
 
     def _batched_call(self, dataset: Dataset) -> Dataset: