Fix CI issues.

zh-plus · Mar 29, 2024 · 02b6b22 · 02b6b22
1 parent c043545
commit 02b6b22
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/openlrc/utils.py b/openlrc/utils.py
@@ -1,7 +1,6 @@
 #  Copyright (C) 2024. Hao Zheng
 #  All rights reserved.
 
-import gc
 import re
 import subprocess
 import time
@@ -74,7 +73,6 @@ def get_audio_duration(path: Union[str, Path]) -> float:
 
 
 def release_memory(model: torch.nn.Module) -> None:
-    gc.collect()
     torch.cuda.empty_cache()
     del model
 

diff --git a/tests/test_chatbot.py b/tests/test_chatbot.py
@@ -10,19 +10,23 @@
 from openlrc.chatbot import GPTBot, ClaudeBot
 
 
-class OpenAIUsage(BaseModel):
+class Usage(BaseModel):
+    pass
+
+
+class OpenAIUsage(Usage):
     prompt_tokens: int
     completion_tokens: int
     total_tokens: int
 
 
-class AnthropicUsage(BaseModel):
+class AnthropicUsage(Usage):
     input_tokens: int
     output_tokens: int
 
 
 class OpenAIResponse(BaseModel):
-    usage: Union[OpenAIUsage | AnthropicUsage]
+    usage: Union[Usage]
 
 
 class TestGPTBot(unittest.TestCase):