add content message field to model quality metrics

api/app/models/metrics/model_quality_dto.py

@@ -199,6 +199,7 @@ class TokenProb(BaseModel):
 
 class TokenData(BaseModel):
     id: str
+    message_content: str
     probs: List[TokenProb]
 
 

api/tests/commons/db_mock.py

@@ -550,6 +550,7 @@ def get_sample_completion_dataset(
     'tokens': [
         {
             'id': 'chatcmpl',
+            'message_content': 'Sky is blue.',
             'probs': [
                 {'prob': 0.27718424797058105, 'token': 'Sky'},
                 {'prob': 0.8951022028923035, 'token': ' is'},

spark/jobs/metrics/completion_metrics.py

@@ -38,20 +38,25 @@ def remove_columns(df: DataFrame) -> DataFrame:
         return df
 
     def compute_prob(self, df: DataFrame):
-        df = df.select(F.explode("choices").alias("element"), F.col("id"))
         df = df.select(
-            F.col("id"), F.explode("element.logprobs.content").alias("content")
+            F.explode("choices").alias("element"),
+            F.col("id"),
         )
-        df = df.select("id", "content.logprob", "content.token").withColumn(
-            "prob", self.compute_probability_udf("logprob")
+        df = df.select(
+            F.col("id"),
+            F.col("element.message.content").alias("message_content"),
+            F.explode("element.logprobs.content").alias("content"),
         )
+        df = df.select(
+            "id", "message_content", "content.logprob", "content.token"
+        ).withColumn("prob", self.compute_probability_udf("logprob"))
         return df
 
     def extract_metrics(self, df: DataFrame) -> CompletionMetricsModel:
         df = self.remove_columns(df)
         df = self.compute_prob(df)
         df_prob = df.drop("logprob")
-        df_prob = df_prob.groupBy("id").agg(
+        df_prob = df_prob.groupBy("id", "message_content").agg(
             F.collect_list(F.struct("token", "prob")).alias("probs")
         )
         df_mean_values = df.groupBy("id").agg(
@@ -66,9 +71,11 @@ def extract_metrics(self, df: DataFrame) -> CompletionMetricsModel:
             F.mean("prob_per_phrase").alias("prob_tot_mean"),
             F.mean("perplex_per_phrase").alias("perplex_tot_mean"),
         )
+        df_prob = df_prob.orderBy("id")
         tokens = [
             {
                 "id": row["id"],
+                "message_content": row["message_content"],
                 "probs": [
                     {"token": prob["token"], "prob": prob["prob"]}
                     for prob in row["probs"]

spark/jobs/models/completion_dataset.py

@@ -9,6 +9,7 @@ class Prob(BaseModel):
 
 class Probs(BaseModel):
     id: str
+    message_content: str
     probs: List[Prob]
 
     model_config = ConfigDict(ser_json_inf_nan="null")

spark/tests/completion_metrics_test.py

@@ -25,7 +25,7 @@ def test_compute_prob(spark_fixture, input_file):
     completion_metrics_service = CompletionMetrics()
     df = completion_metrics_service.remove_columns(input_file)
     df = completion_metrics_service.compute_prob(df)
-    assert {"id", "logprob", "token", "prob"} == set(df.columns)
+    assert {"id", "logprob", "message_content", "token", "prob"} == set(df.columns)
     assert not df.rdd.isEmpty()
 
 

spark/tests/results/completion_metrics_results.py

@@ -2,6 +2,7 @@
     "tokens": [
         {
             "id": "chatcmpl-AcWID2SsE5iuK6z5AhNCKv3WUcCxN",
+            "message_content": "Sure, go ahead. What's up?",
             "probs": [
                 {"token": "Sure", "prob": 0.541987419128418},
                 {"token": ",", "prob": 0.9025230407714844},
@@ -15,6 +16,7 @@
         },
         {
             "id": "chatcmpl-AcYMMPLnpkksCdLze3M8nnqQbfqVG",
+            "message_content": "Certainly! Just let me know how.",
             "probs": [
                 {"token": "Certainly", "prob": 0.022015240043401718},
                 {"token": "!", "prob": 0.8896080851554871},