Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(llm): add support for fine-tuned OpenAI models #682

Merged
merged 1 commit into from
Oct 24, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions pandasai/helpers/openai_info.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,10 @@
"gpt-35-turbo-16k-0613-completion": 0.004,
# Others
"text-davinci-003": 0.02,
# Fine-tuned input
"gpt-3.5-turbo-0613-finetuned": 0.012,
# Fine-tuned output
"gpt-3.5-turbo-0613-finetuned-completion": 0.016,
}


Expand All @@ -62,10 +66,13 @@ def get_openai_token_cost_for_model(
float: Cost in USD.
"""
model_name = model_name.lower()
if "ft:" in model_name:
model_name = model_name.split(":")[1] + "-finetuned"
if is_completion and (
model_name.startswith("gpt-4")
or model_name.startswith("gpt-3.5")
or model_name.startswith("gpt-35")
or "finetuned" in model_name
):
# The cost of completion token is different from
# the cost of prompt tokens.
Expand Down
10 changes: 8 additions & 2 deletions pandasai/llm/openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,9 +102,15 @@ def call(self, instruction: AbstractPrompt, suffix: str = "") -> str:
"""
self.last_prompt = instruction.to_string() + suffix

if self.model in self._supported_chat_models:
if "ft:" in self.model:
# extract "standard" model name from fine-tuned model
model_name = self.model.split(":")[1]
else:
model_name = self.model
Comment on lines +105 to +109
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The code checks if the model name contains "ft:" and extracts the base model name accordingly. This is a good approach to handle fine-tuned models. However, it assumes that the model name after "ft:" is always valid. It would be better to add a check to ensure that the extracted model name is not empty or invalid. This will prevent potential errors when calling the OpenAI API.

        if "ft:" in self.model:
            # extract "standard" model name from fine-tuned model
            model_name = self.model.split(":")[1]
+            if not model_name:
+                raise ValueError("Invalid fine-tuned model name.")
        else:
            model_name = self.model
Committable suggestion (Beta)
Suggested change
if "ft:" in self.model:
# extract "standard" model name from fine-tuned model
model_name = self.model.split(":")[1]
else:
model_name = self.model
if "ft:" in self.model:
# extract "standard" model name from fine-tuned model
model_name = self.model.split(":")[1]
if not model_name:
raise ValueError("Invalid fine-tuned model name.")
else:
model_name = self.model


if model_name in self._supported_chat_models:
response = self.chat_completion(self.last_prompt)
elif self.model in self._supported_completion_models:
elif model_name in self._supported_completion_models:
response = self.completion(self.last_prompt)
else:
raise UnsupportedModelError(self.model)
Expand Down
7 changes: 7 additions & 0 deletions tests/llms/test_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,3 +130,10 @@ def test_call_supported_chat_model(self, mocker, prompt):

result = openai.call(instruction=prompt)
assert result == "response"

def test_call_finetuned_model(self, mocker, prompt):
openai = OpenAI(api_token="test", model="ft:gpt-3.5-turbo:my-org:custom_suffix:id")
mocker.patch.object(openai, "chat_completion", return_value="response")

result = openai.call(instruction=prompt)
assert result == "response"
Comment on lines +134 to +139
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new test case test_call_finetuned_model is testing the call method with a fine-tuned model. It's good to see that the test is mocking the chat_completion method and asserting the expected response. However, it would be beneficial to also verify that the chat_completion method is called with the correct arguments, especially the model name. This will ensure that the call method is correctly extracting the "standard" model name from the fine-tuned model and passing it to the chat_completion method.

-        mocker.patch.object(openai, "chat_completion", return_value="response")
+        chat_completion_mock = mocker.patch.object(openai, "chat_completion", return_value="response")
-        assert result == "response"
+        assert result == "response"
+        chat_completion_mock.assert_called_once_with(instruction=prompt, model="gpt-3.5-turbo")
Committable suggestion (Beta)
Suggested change
def test_call_finetuned_model(self, mocker, prompt):
openai = OpenAI(api_token="test", model="ft:gpt-3.5-turbo:my-org:custom_suffix:id")
mocker.patch.object(openai, "chat_completion", return_value="response")
result = openai.call(instruction=prompt)
assert result == "response"
def test_call_finetuned_model(self, mocker, prompt):
openai = OpenAI(api_token="test", model="ft:gpt-3.5-turbo:my-org:custom_suffix:id")
chat_completion_mock = mocker.patch.object(openai, "chat_completion", return_value="response")
result = openai.call(instruction=prompt)
assert result == "response"
chat_completion_mock.assert_called_once_with(instruction=prompt, model="gpt-3.5-turbo")