damn it i forgot to run lint

Author: lisadunlap

fastchat/serve/monitor/classify/category.py

@@ -79,7 +79,7 @@ def get_score(self, judgment):
 
     def pre_process(self, prompt):
         conv = [{"role": "system", "content": self.sys_prompt}]
-        conv.append({"role": "user", "content": prompt['prompt']})
+        conv.append({"role": "user", "content": prompt["prompt"]})
         return conv
 
     def post_process(self, judgment):
@@ -106,7 +106,7 @@ def get_score(self, judgment):
             return None
 
     def pre_process(self, prompt):
-        args = {"PROMPT": prompt['prompt']}
+        args = {"PROMPT": prompt["prompt"]}
         conv = [
             {"role": "system", "content": self.system_prompt},
             {"role": "user", "content": self.prompt_template.format(**args)},
@@ -140,7 +140,7 @@ def get_score(self, judgment):
             return None
 
     def pre_process(self, prompt):
-        args = {"PROMPT": prompt['prompt']}
+        args = {"PROMPT": prompt["prompt"]}
         conv = [
             {"role": "system", "content": self.system_prompt},
             {"role": "user", "content": self.prompt_template.format(**args)},
@@ -177,7 +177,7 @@ def get_score(self, judgment):
             return None
 
     def pre_process(self, prompt):
-        args = {"PROMPT": prompt['prompt']}
+        args = {"PROMPT": prompt["prompt"]}
         conv = [
             {"role": "system", "content": self.system_prompt},
             {"role": "user", "content": self.prompt_template.format(**args)},
@@ -188,20 +188,19 @@ def post_process(self, judgment):
         score = self.get_score(judgment=judgment)
         bool_score = bool(score == "yes") if score else False
         return {"creative_writing": bool_score, "score": score}
-    
+
+
 #####################
 # Vision Categories #
 #####################
 class CategoryCaptioning(Category):
-
     def __init__(self):
         super().__init__()
         self.name_tag = "captioning_v0.1"
         self.pattern = re.compile(r"<decision>(\w+)<\/decision>")
         self.system_prompt = "You are tasked with determining if a given VQA question is a captioning question. A captioning question asks for a general, overall description of the entire image. It must be a single, open-ended query that does NOT ask about particular objects, people, or parts of the image, nor require interpretation beyond a broad description of what is visually present. Examples include 'What is happening in this image?', 'Describe this picture.', 'Explain', etc. An example of a non-captioning question is 'Describe what is funny in this picture.' because it asks for a specific interpretation of the image content. \n\nOutput your verdict in the following format:<decision>\n[yes/no]\n</decision>. Do NOT explain."
         self.prompt_template = "<user_prompt>\n{PROMPT}\n</user_prompt>"
 
-    
     def get_score(self, judgment):
         matches = self.pattern.findall(judgment.replace("\n", "").lower())
         matches = [m for m in matches if m != ""]
@@ -212,7 +211,7 @@ def get_score(self, judgment):
         else:
             return None
 
-    def pre_process(self, prompt, api_type='openai'):
+    def pre_process(self, prompt, api_type="openai"):
         args = {"PROMPT": prompt["prompt"]}
         conv = [
             {"role": "system", "content": self.system_prompt},
@@ -223,7 +222,8 @@ def pre_process(self, prompt, api_type='openai'):
     def post_process(self, judgment):
         score = self.get_score(judgment=judgment)
         return {"captioning": bool(score == "yes") if score else False}
-    
+
+
 class CategoryCreativeWritingVision(Category):
     def __init__(self):
         super().__init__()
@@ -248,8 +248,8 @@ def get_score(self, judgment):
         else:
             return None
 
-    def pre_process(self, prompt, api_type='openai'):
-        args = {"PROMPT": prompt['prompt']}
+    def pre_process(self, prompt, api_type="openai"):
+        args = {"PROMPT": prompt["prompt"]}
         conv = [
             {"role": "system", "content": self.system_prompt},
             {"role": "user", "content": self.prompt_template.format(**args)},
@@ -260,16 +260,16 @@ def post_process(self, judgment):
         score = self.get_score(judgment=judgment)
         bool_score = bool(score == "yes") if score else False
         return {"creative_writing": bool_score, "score": score}
-    
-class CategoryEntityRecognition(Category):
 
+
+class CategoryEntityRecognition(Category):
     def __init__(self):
         super().__init__()
         self.name_tag = "entity_recognition_v0.1"
         self.pattern = re.compile(r"<decision>(\w+)<\/decision>")
         self.system_prompt = "You are tasked with determining if a given VQA question is an entity recognition question. An entity recognition question asks for the identification of specific objects or people in the image. This does NOT include questions that ask for a general description of the image, questions that only ask for object counts, or questions that only require reading text in the image.\n\nOutput your verdict in the following format:<decision>\n[yes/no]\n</decision>. Do NOT explain."
         self.prompt_template = "<user_prompt>\n{PROMPT}\n</user_prompt>"
-        
+
     def get_score(self, judgment):
         matches = self.pattern.findall(judgment.replace("\n", "").lower())
         matches = [m for m in matches if m != ""]
@@ -280,7 +280,7 @@ def get_score(self, judgment):
         else:
             return None
 
-    def pre_process(self, prompt, api_type='openai'):
+    def pre_process(self, prompt, api_type="openai"):
         args = {"PROMPT": prompt["prompt"]}
         conv = [
             {"role": "system", "content": self.system_prompt},
@@ -291,19 +291,22 @@ def pre_process(self, prompt, api_type='openai'):
     def post_process(self, judgment):
         score = self.get_score(judgment=judgment)
         return {"entity_recognition": bool(score == "yes") if score else False}
-    
+
+
 import base64
 import io
 from PIL import Image
+
+
 def pil_to_base64(image_path):
     image = Image.open(image_path)
     buffered = io.BytesIO()
     image.save(buffered, format="PNG")
     img_str = base64.b64encode(buffered.getvalue()).decode()
     return img_str
 
-class CategoryOpticalCharacterRecognition(Category):
 
+class CategoryOpticalCharacterRecognition(Category):
     def __init__(self):
         super().__init__()
         self.name_tag = "ocr_v0.1"
@@ -321,21 +324,23 @@ def get_score(self, judgment):
         else:
             return None
 
-    def pre_process(self, prompt, api_type='openai'):
+    def pre_process(self, prompt, api_type="openai"):
         args = {"PROMPT": prompt["prompt"]}
         base64_image = pil_to_base64(prompt["image_path"])
-        if api_type == 'anthropic':
+        if api_type == "anthropic":
             conv = [
                 {"role": "system", "content": self.system_prompt},
                 {
                     "role": "user",
                     "content": [
-                            {
+                        {
                             "type": "image",
                             "source": {
                                 "type": "base64",
                                 "media_type": "image/jpeg",
-                                "data": base64.b64encode(prompt["image_path"].content).decode("utf-8"),
+                                "data": base64.b64encode(
+                                    prompt["image_path"].content
+                                ).decode("utf-8"),
                             },
                         },
                         {"type": "text", "text": self.prompt_template.format(**args)},
@@ -363,9 +368,9 @@ def pre_process(self, prompt, api_type='openai'):
     def post_process(self, judgment):
         score = self.get_score(judgment=judgment)
         return {"ocr": bool(score == "yes") if score else False}
-    
-class CategoryHumor(Category):
 
+
+class CategoryHumor(Category):
     def __init__(self):
         super().__init__()
         self.name_tag = "humor_v0.1"
@@ -382,17 +387,17 @@ def get_score(self, judgment):
             return matches[0]
         else:
             return None
-        
-    def pre_process(self, prompt, api_type='openai'):
+
+    def pre_process(self, prompt, api_type="openai"):
         args = {"PROMPT": prompt["prompt"]}
         base64_image = pil_to_base64(prompt["image_path"])
-        if api_type == 'anthropic':
+        if api_type == "anthropic":
             conv = [
                 {"role": "system", "content": self.system_prompt},
                 {
                     "role": "user",
                     "content": [
-                            {
+                        {
                             "type": "image",
                             "source": {
                                 "type": "base64",
@@ -421,14 +426,16 @@ def pre_process(self, prompt, api_type='openai'):
                 },
             ]
         return conv
-    
+
     def post_process(self, judgment):
         score = self.get_score(judgment=judgment)
         return {"humor": bool(score == "yes") if score else False}
-    
+
+
 import os
-class CategoryHomework(Category):
 
+
+class CategoryHomework(Category):
     def __init__(self):
         super().__init__()
         self.name_tag = "homework_v0.1"
@@ -449,21 +456,21 @@ def get_score(self, judgment):
             return matches[0]
         else:
             return None
-        
-    def pre_process(self, prompt, api_type='openai'):
+
+    def pre_process(self, prompt, api_type="openai"):
         base64_image = pil_to_base64(prompt["image_path"])
 
         # Open the local image file in binary mode and encode it as base64
         assert os.path.exists(prompt["image_path"])
         with open(prompt["image_path"], "rb") as image_file:
             image_data = base64.b64encode(image_file.read()).decode("utf-8")
-        if api_type == 'anthropic':
+        if api_type == "anthropic":
             conv = [
                 {"role": "system", "content": self.system_prompt},
                 {
                     "role": "user",
                     "content": [
-                            {
+                        {
                             "type": "image",
                             "source": {
                                 "type": "base64",
@@ -492,13 +499,13 @@ def pre_process(self, prompt, api_type='openai'):
                 },
             ]
         return conv
-    
+
     def post_process(self, judgment):
         score = self.get_score(judgment=judgment)
         return {"homework": bool(score == "yes") if score else False}
-    
-class CategoryDiagram(Category):
 
+
+class CategoryDiagram(Category):
     def __init__(self):
         super().__init__()
         self.name_tag = "diagram_v0.1"
@@ -523,21 +530,21 @@ def get_score(self, judgment):
             return matches[0]
         else:
             return None
-        
-    def pre_process(self, prompt, api_type='openai'):
+
+    def pre_process(self, prompt, api_type="openai"):
         base64_image = pil_to_base64(prompt["image_path"])
 
         # Open the local image file in binary mode and encode it as base64
         assert os.path.exists(prompt["image_path"])
         with open(prompt["image_path"], "rb") as image_file:
             image_data = base64.b64encode(image_file.read()).decode("utf-8")
-        if api_type == 'anthropic':
+        if api_type == "anthropic":
             conv = [
                 {"role": "system", "content": self.system_prompt},
                 {
                     "role": "user",
                     "content": [
-                            {
+                        {
                             "type": "image",
                             "source": {
                                 "type": "base64",
@@ -566,7 +573,7 @@ def pre_process(self, prompt, api_type='openai'):
                 },
             ]
         return conv
-    
+
     def post_process(self, judgment):
         score = self.get_score(judgment=judgment)
         return {"diagram": bool(score == "yes") if score else False}

fastchat/serve/monitor/classify/label.py

@@ -88,6 +88,7 @@ def chat_completion_openai(model, messages, temperature, max_tokens, api_dict=No
 
     return output
 
+
 def chat_completion_anthropic(model, messages, temperature, max_tokens, api_dict=None):
     import anthropic
 
@@ -111,7 +112,7 @@ def chat_completion_anthropic(model, messages, temperature, max_tokens, api_dict
                 stop_sequences=[anthropic.HUMAN_PROMPT],
                 max_tokens=max_tokens,
                 temperature=temperature,
-                system=sys_msg
+                system=sys_msg,
             )
             output = response.content[0].text
             break
@@ -120,7 +121,10 @@ def chat_completion_anthropic(model, messages, temperature, max_tokens, api_dict
             time.sleep(API_RETRY_SLEEP)
     return output
 
-def chat_completion_gemini(model, messages, temperature, max_tokens, api_dict=None, image_path=None):
+
+def chat_completion_gemini(
+    model, messages, temperature, max_tokens, api_dict=None, image_path=None
+):
     import google
     import google.generativeai as genai
     from google.generativeai.types import HarmCategory, HarmBlockThreshold
@@ -139,9 +143,9 @@ def chat_completion_gemini(model, messages, temperature, max_tokens, api_dict=No
 
     prompt = messages[0]["content"]
     if type(prompt) == list:
-        prompt = [prompt[0]['text'], Image.open(image_path).convert('RGB')]
+        prompt = [prompt[0]["text"], Image.open(image_path).convert("RGB")]
 
-    safety_settings={
+    safety_settings = {
         HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
         HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
         HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
@@ -156,7 +160,9 @@ def chat_completion_gemini(model, messages, temperature, max_tokens, api_dict=No
             gemini.temperature = temperature
             response = gemini.generate_content(prompt, safety_settings=safety_settings)
             if response.candidates[0].finish_reason != 1:
-                print(f"Gemini did not finish generating content: {response.candidates[0].finish_reason}")
+                print(
+                    f"Gemini did not finish generating content: {response.candidates[0].finish_reason}"
+                )
                 output = "Gemini did not finish generating content"
             else:
                 output = response.text
@@ -215,7 +221,7 @@ def get_answer(
                 temperature=temperature,
                 max_tokens=max_tokens,
                 api_dict=api_dict,
-                image_path=question.get("image_path")
+                image_path=question.get("image_path"),
             )
         else:
             raise ValueError(f"api_type {api_type} not supported")
@@ -309,7 +315,9 @@ def find_required_tasks(row):
         input_data["image_hash"] = input_data.conversation_a.map(
             lambda convo: convo[0]["content"][1][0]
         )
-        input_data["image_path"] = input_data.image_hash.map(lambda x: f"{config['image_dir']}/{x}.png")
+        input_data["image_path"] = input_data.image_hash.map(
+            lambda x: f"{config['image_dir']}/{x}.png"
+        )
 
     if config["cache_file"]:
         print("loading cache data")
@@ -360,11 +368,15 @@ def find_required_tasks(row):
 
     if args.vision:
         not_labeled["prompt"] = not_labeled.conversation_a.map(
-            lambda convo: "\n".join([convo[i]["content"][0] for i in range(0, len(convo), 2)])
+            lambda convo: "\n".join(
+                [convo[i]["content"][0] for i in range(0, len(convo), 2)]
+            )
         )
     else:
         not_labeled["prompt"] = not_labeled.conversation_a.map(
-            lambda convo: "\n".join([convo[i]["content"] for i in range(0, len(convo), 2)])
+            lambda convo: "\n".join(
+                [convo[i]["content"] for i in range(0, len(convo), 2)]
+            )
         )
     not_labeled["prompt"] = not_labeled.prompt.map(lambda x: x[:12500])