Pass extra arguments from DAILA (#9)

Co-authored-by: Edward J Schwartz <moo.github.z.edmcman@xoxy.net>
binsync · Feb 23, 2024 · e8c8c15 · e8c8c15
1 parent a154a77
commit e8c8c15
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/varbert/api.py b/varbert/api.py
@@ -98,13 +98,13 @@ def predict_variable_names(
         return orig_name_2_popular_name, renamed_code
 
     @AIAPI.requires_function
-    def query_model(self, *args, function=None, dec_text=None, use_dec=True, **kwargs) -> Tuple[Dict[str, str], str]:
+    def query_model(self, function=None, dec_text=None, use_dec=True, *args, **kwargs) -> Tuple[Dict[str, str], str]:
         """
         Standardized function for querying the model in DAILA interfaces.
         """
-        old_to_new_vars, renamed_code = self.predict_variable_names(function=function, decompilation_text=dec_text, use_decompiler=use_dec)
+        old_to_new_vars, renamed_code = self.predict_variable_names(function=function, decompilation_text=dec_text, use_decompiler=use_dec, *args, **kwargs)
         if use_dec and old_to_new_vars:
-            self._dec_interface.rename_local_variables_by_names(function, old_to_new_vars)
+            self._dec_interface.rename_local_variables_by_names(function, old_to_new_vars, *args, **kwargs)
 
         return old_to_new_vars, renamed_code