feat: completion model prompts

Co1lin · Nov 14, 2024 · e484827 · e484827
1 parent f4af08d
commit e484827
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 20 deletions.
diff --git a/cweval/ai/google.py b/cweval/ai/google.py
@@ -36,26 +36,30 @@ def send_message(
         all_kwargs.update(kwargs)
         # openai kwargs to google gemini ones
         all_kwargs['max_output_tokens'] = all_kwargs.pop('max_completion_tokens', None)
-        all_kwargs['candidate_count'] = all_kwargs.pop('n', None)
-
-        resp = self._model.generate_content(
-            contents=contents,
-            generation_config={
-                **all_kwargs,
-                # 'stop_sequences': stop,
-                # 'temperature': 0.6,
-                # 'top_p': 0.9,
-                # 'top_k': 1,
-            },
-            safety_settings={
-                'sex': 'block_none',
-                'hate': 'block_none',
-                'harassment': 'block_none',
-                'danger': 'block_none',
-            },
-            request_options={"timeout": 2000},
-        )
-        return [' '.join([p.text for p in c.content.parts]) for c in resp.candidates]
+        num_samples = all_kwargs.pop('n', None)
+        resps = []
+        for _ in range(num_samples):
+            resp = self._model.generate_content(
+                contents=contents,
+                generation_config={
+                    **all_kwargs,
+                    # 'stop_sequences': stop,
+                    # 'temperature': 0.6,
+                    # 'top_p': 0.9,
+                    # 'top_k': 1,
+                },
+                safety_settings={
+                    'sex': 'block_none',
+                    'hate': 'block_none',
+                    'harassment': 'block_none',
+                    'danger': 'block_none',
+                },
+                request_options={"timeout": 2000},
+            )
+            resps.extend(
+                [' '.join([p.text for p in c.content.parts]) for c in resp.candidates]
+            )
+        return resps
 
     def chat_send_message(
         self,

diff --git a/cweval/ppt/__init__.py b/cweval/ppt/__init__.py
@@ -65,10 +65,27 @@ class SecurePrompt(DirectPrompt):
     }
 
 
+class DirectCompletionPrompt(DirectPrompt):
+    PPT = '''
+The following code always exit with code 0. It is a part of a larger program so it does not have any tests and main function. {lang_instr}It is wrapped by triple backticks in a single code block.
+
+```{lang}
+{code_prompt}
+'''
+    LANG_INSTR = {
+        lang: instr.replace(
+            'You should implement the function', 'The function is implemented'
+        )
+        for lang, instr in DirectPrompt.LANG_INSTR.items()
+    }
+
+
 def make_prompt(ppt: str) -> Prompt:
     if ppt == 'direct':
         return DirectPrompt
     elif ppt == 'secure':
         return SecurePrompt
+    elif ppt == 'compl':
+        return DirectCompletionPrompt
     else:
         raise NotImplementedError(f'Unknown prompt type: {ppt}')