Skip to content
Snippets Groups Projects
Commit 21bd06d3 authored by Silin's avatar Silin
Browse files

Update local_evaluation_with_api.py, add multi-message prompt.

parent 2dc006e1
No related branches found
No related tags found
No related merge requests found
......@@ -38,12 +38,20 @@ class LLM_API:
self.model = "gpt-3.5-turbo-1106"
def api_call(self, prompt, max_tokens):
""" Simple single prompt api call """
response = self.client.chat.completions.create(
model=self.model,
messages=[{"role": "user", "content": prompt}],
max_tokens=max_tokens,
)
if isinstance(prompt, str): # Single-message prompt
response = self.client.chat.completions.create(
model=self.model,
messages=[{"role": "user", "content": prompt}],
max_tokens=max_tokens,
)
elif isinstance(prompt, list): # Multi-message prompt
response = self.client.chat.completions.create(
model=self.model,
messages=prompt,
max_tokens=max_tokens,
)
else:
raise TypeError
response_text = response.choices[0].message.content
input_tokens = response.usage.prompt_tokens
output_tokens = response.usage.completion_tokens
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment