adding utilities for generating prompts

prompting/promptingutils.py

@@ -0,0 +1,190 @@
+import os, sys
+import pandas as pd
+
+
+TOP_COMMENTARIES_DIR = os.path.join(os.curdir, "top_commentaries")
+
+
+ARTEMIS_EMOTIONS = [
+    "amusement",         
+    "anger",              
+    "awe",               
+    "contentment",       
+    "disgust",           
+    "excitement",        
+    "fear",              
+    "sadness",           
+    "something else",    
+]
+
+DEFAULT_N_SAMPLES = 4
+DEFAULT_RANDOM_STATE = 42
+DEFAULT_OBJECT_THRESHOLD = 0.1
+
+
+
+
+
+def fill_extracted_items(
+    colors_list = None,
+    objects_list = None,
+    emotion = None,
+    commentary = None,
+):
+  result = ""
+  if colors_list:
+    result += f"colors: {', '.join(colors_list)}"
+    result += "\n"
+  if objects_list:
+    result += f"objects: {', '.join(objects_list)}"
+    result += "\n"
+  if emotion:
+    result += f"emotion: {emotion}"
+    result += "\n"
+  if commentary:
+    result += f"commentary: {commentary}"
+    result += "\n"
+  else:
+    result += "commentary: "
+
+  return result
+
+
+def load_dataframe(
+    csv_filepath,
+):
+  df = pd.read_csv(csv_filepath, index_col = 0)
+  for stringified_col in  [
+                            "maskrcnn_objects",
+                            "colors",
+                            "clip_recognized_objects", 
+                          ]:
+      df[stringified_col] = df[stringified_col].apply(eval)
+
+  return df
+
+TOP_COMMENTARIES_DFS = {
+    emotion : load_dataframe(os.path.join(TOP_COMMENTARIES_DIR, f"top_{emotion.replace(' ', '_')}.csv"))
+    for emotion in ARTEMIS_EMOTIONS
+    }
+  
+
+def filter_items(
+    items_and_probs_list,
+    items_threshold = DEFAULT_OBJECT_THRESHOLD
+    ):
+    return [item for item, prob in items_and_probs_list if prob > items_threshold]
+
+def get_random_samples_for_emotion(
+    emotion,
+    n_samples = DEFAULT_N_SAMPLES,
+    random_state = DEFAULT_RANDOM_STATE,
+    object_threshold = DEFAULT_OBJECT_THRESHOLD
+):
+    emotion_df = TOP_COMMENTARIES_DFS[emotion]
+    samples = emotion_df.sample(n_samples, random_state = random_state)
+
+
+
+    result = []
+    for _, sample in samples.iterrows():
+      colors_list = sample["colors"]
+      commentary = sample["utterance"]
+      emotion = sample["emotion"]
+
+      objects_and_probs_list = sample["clip_recognized_objects"]
+      objects_list = filter_items(objects_and_probs_list, object_threshold)
+
+      entry = {
+        "colors_list" : colors_list,
+        "objects_list" : objects_list,
+        "emotion" : emotion,
+        "commentary" : commentary,
+      }
+      result.append(entry)
+
+    return result
+
+def get_subprompt_for_emotion(
+    emotion,
+    n_samples = DEFAULT_N_SAMPLES,
+    random_state = DEFAULT_RANDOM_STATE,
+    object_threshold = DEFAULT_OBJECT_THRESHOLD,
+):
+  random_samples = get_random_samples_for_emotion(
+      emotion = emotion,
+      n_samples = n_samples,
+      random_state = random_state,
+      object_threshold=object_threshold,
+  )
+  subprompt = [
+        fill_extracted_items(**entry) for entry in random_samples    
+  ]
+  subprompt = "\n".join(subprompt)
+
+  return subprompt
+
+
+def get_subprompt_with_examples(
+    n_samples_per_emotion = DEFAULT_N_SAMPLES,
+    random_state = DEFAULT_RANDOM_STATE,
+    object_threshold = DEFAULT_OBJECT_THRESHOLD,
+):
+
+    examples = [
+        get_subprompt_for_emotion(
+            emotion=emotion,
+            n_samples = n_samples_per_emotion,
+            random_state = random_state,
+            object_threshold = object_threshold
+        )   
+        for emotion in ARTEMIS_EMOTIONS
+    ]
+
+    examples = "\n".join(examples)
+
+    return examples
+
+def get_user_prompt(
+    colors_list,
+    objects_list,
+    emotion,
+    n_samples_per_emotion = DEFAULT_N_SAMPLES,
+    random_state = DEFAULT_RANDOM_STATE,
+    object_threshold = DEFAULT_OBJECT_THRESHOLD,
+):
+  user_prompt= (
+      "You have to write a commentary for an artwork.\n"
+      "To write the commentary, you are given the objects present in the picture, "
+      "the colors present in the picture, and the emotion the picture evokes.\n"
+      "You are first shown several examples, and then have to give your commentary.\n"
+      "First come the examples, and then the objects, colors, and emotion you will have to use for your commentary.\n"
+      "Avoid explicitly mentioning the objects, or colors, or emotion, if it sounds more natural.\n"
+      "Only write the commentary.\n"
+      "\n"
+      "EXAMPLES:"
+      "\n\n"
+      "{examples}"
+      "\n"
+      "Now, write your personal opinion about the picture."
+      "\n"
+      "{image_subprompt}"
+  )
+
+  examples = get_subprompt_with_examples(
+      n_samples_per_emotion = n_samples_per_emotion,  
+      random_state = random_state,
+      object_threshold=object_threshold,
+  )
+
+  image_subprompt = fill_extracted_items(
+    colors_list = colors_list,
+    objects_list = objects_list,
+    emotion = emotion,
+    commentary = None,
+  )
+
+
+  result = user_prompt.format(examples = examples, image_subprompt = image_subprompt)
+
+  return result