Add instruct prompt

dotsd · Dec 13, 2023 · 922e4e0 · 922e4e0
1 parent eb3d6c2
commit 922e4e0
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/main.py b/main.py
@@ -111,12 +111,14 @@ def generate(prompts: List[str], model: Transformer, tokenizer: Tokenizer, *, ma
     return generated_words, logprobs
 
 
-def interactive(model_path: str, max_tokens: int = 35, temperature: float = 0.7):
+def interactive(model_path: str, max_tokens: int = 35, temperature: float = 0.7, instruct: bool = False):
     tokenizer = Tokenizer(str(Path(model_path) / "tokenizer.model"))
     transformer = Transformer.from_folder(Path(model_path), max_batch_size=3)
 
     while True:
         prompt = input("Prompt: ")
+        if instruct:
+            prompt = f"[INST] {prompt} [/INST]"
         res, _logprobs = generate(
             [prompt],
             transformer,