Fixing new default scoring

flowersteam · Nov 21, 2023 · 4b36c0c · 4b36c0c
1 parent 744ffa3
commit 4b36c0c
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 5 deletions.
diff --git a/lamorel/src/lamorel/server/llms/module_functions/__init__.py b/lamorel/src/lamorel/server/llms/module_functions/__init__.py
@@ -1,2 +1,2 @@
 from .base_module_function import BaseModuleFunction
-from .score_module_function import ScoreModuleFunction
+from .score_module_function import LogScoringModuleFn
diff --git a/lamorel/src/lamorel/server/llms/module_functions/score_module_function.py b/lamorel/src/lamorel/server/llms/module_functions/score_module_function.py
@@ -2,10 +2,10 @@
 import torch
 
 class LogScoringModuleFn(BaseModuleFunction):
-    def __init__(self, model_type, pre_encoded_input):
+    def __init__(self, pade_token, model_type, pre_encoded_input):
         super().__init__()
+        self._pad_token = pade_token
         self._model_type = model_type
-        self._pad_token = 0
         self._pre_encoded_input = pre_encoded_input
 
     def initialize(self):

diff --git a/lamorel/src/lamorel/server/server.py b/lamorel/src/lamorel/server/server.py
@@ -9,7 +9,7 @@
 
 from .llms import HF_LLM
 from .llms.updaters import BaseUpdater
-from .llms.module_functions import BaseModuleFunction, ScoreModuleFunction
+from .llms.module_functions import BaseModuleFunction, LogScoringModuleFn
 from .dispatcher import Dispatcher
 from .utils import InstructionsEnum
 
@@ -41,7 +41,8 @@ def __init__(self, config, llm_index, llm_group, llm_master, rl_llm_group, rl_ll
         self._dispatcher = Dispatcher(self._llm_group, self._rl_llm_group_size - 1, self._llm_group_size,
                                       self._is_main_server, self._master_server_rank, self._index)
 
-        custom_module_functions["__score"] = ScoreModuleFunction(self._model.pad_token, config.llm_args.model_type)
+        custom_module_functions["__score"] = LogScoringModuleFn(self._model.pad_token, config.llm_args.model_type,
+                                                                config.llm_args.pre_encode_inputs)
         for k, _fn in custom_module_functions.items():
             assert isinstance(_fn, BaseModuleFunction)
             _fn.device = self._model.device