fix types

llama_lora/ui/css_styles.py

@@ -1,4 +1,6 @@
-css_styles = []
+from typing import List
+
+css_styles: List[str] = []
 
 
 def get_css_styles():

llama_lora/ui/inference_ui.py

@@ -556,9 +556,10 @@ def inference_ui():
                             elem_id="inference_inference_raw_output_accordion"
                     ) as raw_output_group:
                         inference_raw_output = gr.Code(
-                            label="Raw Output",
-                            show_label=False,
+                            # label="Raw Output",
+                            label="Tensor",
                             language="json",
+                            lines=8,
                             interactive=False,
                             elem_id="inference_raw_output")
 

llama_lora/ui/tokenizer_ui.py

@@ -46,6 +46,7 @@ def tokenizer_ui():
                 encoded_tokens = gr.Code(
                     label="Encoded Tokens (JSON)",
                     language="json",
+                    lines=10,
                     value=sample_encoded_tokens_value,
                     elem_id="tokenizer_encoded_tokens_input_textbox")
                 decode_btn = gr.Button("Decode ➡️")
@@ -54,6 +55,7 @@ def tokenizer_ui():
             with gr.Column():
                 decoded_tokens = gr.Code(
                     label="Decoded Tokens",
+                    lines=10,
                     value=sample_decoded_text_value,
                     elem_id="tokenizer_decoded_text_input_textbox")
                 encode_btn = gr.Button("⬅️ Encode")

llama_lora/utils/prompter.py

@@ -7,7 +7,7 @@ import json
 import os.path as osp
 import importlib
 import itertools
-from typing import Union, List
+from typing import Union, List, Dict
 
 from ..config import Config
 from ..globals import Global
@@ -38,9 +38,10 @@ class Prompter(object):
             raise ValueError(f"Can't read {file_path}")
 
         if ext == ".py":
-            template_module_spec = importlib.util.spec_from_file_location(
+            importlib_util = importlib.util  # type: ignore
+            template_module_spec = importlib_util.spec_from_file_location(
                 "template_module", file_path)
-            template_module = importlib.util.module_from_spec(
+            template_module = importlib_util.module_from_spec(
                 template_module_spec)
             template_module_spec.loader.exec_module(template_module)
             self.template_module = template_module
@@ -67,7 +68,7 @@ class Prompter(object):
 
     def generate_prompt(
         self,
-        variables: List[Union[None, str]] = [],
+        variables: Union[Dict[str, str], List[Union[None, str]]] = [],
         # instruction: str,
         # input: Union[None, str] = None,
         label: Union[None, str] = None,
@@ -75,10 +76,14 @@ class Prompter(object):
         if self.template_name == "None":
             if type(variables) == list:
                 res = get_val(variables, 0, "")
-            else:
+            elif type(variables) == dict:
                 res = variables.get("prompt", "")
+            else:
+                raise ValueError(f"Invalid variables type: {type(variables)}")
         elif "variables" in self.template:
             variable_names = self.template.get("variables")
+            # if type(variable_names) != list:
+            #     raise ValueError(f"Invalid variable_names type {type(variable_names)} defined in template {self.template_name}, expecting list.")
             if self.template_module:
                 if type(variables) == list:
                     variables = {k: v for k, v in zip(