Add llm explanation to output dataframe (#462)

* Add llm explanation to output dataframe

* formatting fixes

---------

Co-authored-by: Rajas Bansal <[email protected]>

docs/guide/accuracy/chain-of-thought.md

@@ -44,19 +44,20 @@ config = {
         "few_shot_selection": "semantic_similarity",
         "few_shot_num": 3,
         "example_template": "Context: {context}\nQuestion: {question}\nAnswer: Let's think step by step.\n{explanation}\n{answer}",
-        "chain_of_thought": true
+        "chain_of_thought": True
     }
 }
 ```
 
 Notice the changes that we have made to the config compared to the config without Chain-of-Thought [here](/guide/tasks/question_answering_task):
 
-* `chain_of_thought` flag - this tells labeling agent to expect an explanation for the answer, in the seed dataset as well as LLM generated responses.  
-* `explanation_column` - this is the column where the explanation for the seed examples will reside.
-* `example_template` - Notice that the template contains contains the explanation column as well. This tells the config where the explanation should be put when using the seed examples. We use the `Let's think step by step` prompt to initiate the chain of thought in the model.
-* `output_guidelines` - We are explicitly prompting the LLM to first output an explanation, and then the final answer.
+- `chain_of_thought` flag - this tells labeling agent to expect an explanation for the answer, in the seed dataset as well as LLM generated responses.
+- `explanation_column` - this is the column where the explanation for the seed examples will reside.
+- `example_template` - Notice that the template contains contains the explanation column as well. This tells the config where the explanation should be put when using the seed examples. We use the `Let's think step by step` prompt to initiate the chain of thought in the model.
+- `output_guidelines` - We are explicitly prompting the LLM to first output an explanation, and then the final answer.
 
 Now, in order to generate explanations for the seed examples, in case they were not manually generated is,
+
 ```py
 from autolabel import LabelingAgent
 agent = LabelingAgent(config)
@@ -68,6 +69,7 @@ Once these explanations are generated, the dataset looks like
 {{ read_csv('docs/assets/squad_with_explanation_preview.csv') }}
 
 Now to generate labels for this dataset, all we have to do is,
+
 ```py
 agent.plan('data/squad_v2_test.csv')
 agent.run('data/squad_v2_test.csv', max_items = 100)

src/autolabel/labeler.py

@@ -259,7 +259,14 @@ def run(
             l.label for l in llm_labels
         ]
         if self.config.confidence():
-            output_df["llm_confidence"] = [l.confidence_score for l in llm_labels]
+            output_df[self.config.task_name() + "_llm_confidence"] = [
+                l.confidence_score for l in llm_labels
+            ]
+
+        if self.config.chain_of_thought():
+            output_df[self.config.task_name() + "_llm_explanation"] = [
+                l.explanation for l in llm_labels
+            ]
 
         # Only save to csv if output_name is provided or dataset is a string
         if not output_name and isinstance(dataset, str):

src/autolabel/schema.py

@@ -80,6 +80,7 @@ class LLMAnnotation(BaseModel):
     confidence_score: Optional[float] = None
     generation_info: Optional[Dict[str, Any]] = None
     raw_response: Optional[str] = ""
+    explanation: Optional[str] = ""
     prompt: Optional[str] = ""
 
 

src/autolabel/tasks/base.py

@@ -69,6 +69,7 @@ def parse_llm_response(
         # This is done to handle the case where the model generates an explanation before generating the label
         if self.config.chain_of_thought():
             try:
+                explanation = response.text.strip().split("\n")[0].strip()
                 completion_text = extract_valid_json_substring(
                     response.text.strip().split("\n")[-1].strip()
                 )
@@ -106,4 +107,5 @@ def parse_llm_response(
             raw_response=response.text,
             prompt=prompt,
             curr_sample=json.dumps(curr_sample),
+            explanation=explanation if self.config.chain_of_thought() else "",
         )