forked from openai/evals
-
Notifications
You must be signed in to change notification settings - Fork 0
/
openai.py
139 lines (119 loc) · 4.83 KB
/
openai.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
from typing import Any, Optional, Union
from evals.api import CompletionFn, CompletionResult
from evals.base import CompletionFnSpec
from evals.prompt.base import (
ChatCompletionPrompt,
CompletionPrompt,
OpenAICreateChatPrompt,
OpenAICreatePrompt,
Prompt,
)
from evals.record import record_sampling
from evals.utils.api_utils import (
openai_chat_completion_create_retrying,
openai_completion_create_retrying,
)
class OpenAIBaseCompletionResult(CompletionResult):
def __init__(self, raw_data: Any, prompt: Any):
self.raw_data = raw_data
self.prompt = prompt
def get_completions(self) -> list[str]:
raise NotImplementedError
class OpenAIChatCompletionResult(OpenAIBaseCompletionResult):
def get_completions(self) -> list[str]:
completions = []
if self.raw_data and "choices" in self.raw_data:
for choice in self.raw_data["choices"]:
if "message" in choice:
completions.append(choice["message"]["content"])
return completions
class OpenAICompletionResult(OpenAIBaseCompletionResult):
def get_completions(self) -> list[str]:
completions = []
if self.raw_data and "choices" in self.raw_data:
for choice in self.raw_data["choices"]:
if "text" in choice:
completions.append(choice["text"])
return completions
class OpenAICompletionFn(CompletionFn):
def __init__(
self,
model: Optional[str] = None,
api_base: Optional[str] = None,
api_key: Optional[str] = None,
n_ctx: Optional[int] = None,
extra_options: Optional[dict] = {},
**kwargs,
):
self.model = model
self.api_base = api_base
self.api_key = api_key
self.n_ctx = n_ctx
self.extra_options = extra_options
def __call__(
self,
prompt: Union[str, OpenAICreateChatPrompt],
**kwargs,
) -> OpenAICompletionResult:
if not isinstance(prompt, Prompt):
assert (
isinstance(prompt, str)
or (isinstance(prompt, list) and all(isinstance(token, int) for token in prompt))
or (isinstance(prompt, list) and all(isinstance(token, str) for token in prompt))
or (isinstance(prompt, list) and all(isinstance(msg, dict) for msg in prompt))
), f"Got type {type(prompt)}, with val {type(prompt[0])} for prompt, expected str or list[int] or list[str] or list[dict[str, str]]"
prompt = CompletionPrompt(
raw_prompt=prompt,
)
openai_create_prompt: OpenAICreatePrompt = prompt.to_formatted_prompt()
result = openai_completion_create_retrying(
model=self.model,
api_base=self.api_base,
api_key=self.api_key,
prompt=openai_create_prompt,
**{**kwargs, **self.extra_options},
)
result = OpenAICompletionResult(raw_data=result, prompt=openai_create_prompt)
record_sampling(prompt=result.prompt, sampled=result.get_completions())
return result
class OpenAIChatCompletionFn(CompletionFnSpec):
def __init__(
self,
model: Optional[str] = None,
api_base: Optional[str] = None,
api_key: Optional[str] = None,
n_ctx: Optional[int] = None,
extra_options: Optional[dict] = {},
**kwargs,
):
self.model = model
self.api_base = api_base
self.api_key = api_key
self.n_ctx = n_ctx
self.extra_options = extra_options
def __call__(
self,
prompt: Union[str, OpenAICreateChatPrompt],
**kwargs,
) -> OpenAIChatCompletionResult:
if not isinstance(prompt, Prompt):
assert (
isinstance(prompt, str)
or (isinstance(prompt, list) and all(isinstance(token, int) for token in prompt))
or (isinstance(prompt, list) and all(isinstance(token, str) for token in prompt))
or (isinstance(prompt, list) and all(isinstance(msg, dict) for msg in prompt))
), f"Got type {type(prompt)}, with val {type(prompt[0])} for prompt, expected str or list[int] or list[str] or list[dict[str, str]]"
prompt = ChatCompletionPrompt(
raw_prompt=prompt,
)
openai_create_prompt: OpenAICreateChatPrompt = prompt.to_formatted_prompt()
result = openai_chat_completion_create_retrying(
model=self.model,
api_base=self.api_base,
api_key=self.api_key,
messages=openai_create_prompt,
**{**kwargs, **self.extra_options},
)
result = OpenAIChatCompletionResult(raw_data=result, prompt=openai_create_prompt)
record_sampling(prompt=result.prompt, sampled=result.get_completions())
return result