From d5d4f47371c8459e67759356803fd50534b12b61 Mon Sep 17 00:00:00 2001
From: Haodong Duan <dhd.efz@gmail.com>
Date: Wed, 9 Aug 2023 12:38:57 +0800
Subject: [PATCH] [API] Refine OpenAI  (#175)

---
 opencompass/models/openai_api.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/opencompass/models/openai_api.py b/opencompass/models/openai_api.py
index 7980221d..5c01e999 100644
--- a/opencompass/models/openai_api.py
+++ b/opencompass/models/openai_api.py
@@ -80,6 +80,7 @@ class OpenAI(BaseAPIModel):
             self.orgs = org
         self.org_ctr = 0
         self.url = openai_api_base
+        self.path = path
 
     def generate(
         self,
@@ -146,9 +147,17 @@ class OpenAI(BaseAPIModel):
                 messages.append(msg)
 
         # max num token for gpt-3.5-turbo is 4097
+        context_window = 4096
+        if '32k' in self.path:
+            context_window = 32768
+        elif '16k' in self.path:
+            context_window = 16384
+        elif 'gpt-4' in self.path:
+            context_window = 8192
+
+        # Hold out 100 tokens due to potential errors in tiktoken calculation
         max_out_len = min(
-            max_out_len,
-            self.max_seq_len - 50 - self.get_token_len(str(input)))
+            max_out_len, context_window - self.get_token_len(str(input)) - 100)
         if max_out_len <= 0:
             return ''