fix test

okdshin · Dec 9, 2023 · db786f7 · db786f7
1 parent 4c0a2be
commit db786f7
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/flatline_lsp.py b/flatline_lsp.py
@@ -257,11 +257,14 @@ def completions(
     document = ls.workspace.get_text_document(params.text_document.uri)
     line_index = params.position.line
     character_index = params.position.character
-    assert lm_for_completion is not None
+    if lm_for_completion is not None:
+        max_context_lines = lm_for_completion.max_context_lines
+    else:
+        max_context_lines = 4
     prompt = "".join(
         list(
             document.lines[
-                max(0, line_index - lm_for_completion.max_context_lines) : line_index
+                max(0, line_index - max_context_lines) : line_index
             ]
         )
         + [document.lines[line_index][:character_index]]

diff --git a/tests/test_language_model_for_completion.py b/tests/test_language_model_for_completion.py
@@ -59,6 +59,10 @@ def test_language_model_for_completion():
     lm_for_completion = flatline_lsp.LanguageModelForCompletion(
         lang_server=server,
         max_new_tokens=32,
+        min_new_tokens=16,
+        max_context_lines=4,
+        max_new_lines=2,
+        max_context_tokens=128,
         tokenizer=tokenizer,
         model=model,
     )