Merge pull request #390 from tomorrmato/add_example_type_hint

added type hint in example code

example_chat_completion.py

@@ -1,11 +1,11 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # This software may be used and distributed according to the terms of the Llama 2 Community License Agreement.
 
-from typing import Optional
+from typing import List, Optional
 
 import fire
 
-from llama import Llama
+from llama import Llama, Dialog
 
 
 def main(
@@ -39,7 +39,7 @@ def main(
         max_batch_size=max_batch_size,
     )
 
-    dialogs = [
+    dialogs: List[Dialog] = [
         [{"role": "user", "content": "what is the recipe of mayonnaise?"}],
         [
             {"role": "user", "content": "I am going to Paris, what should I see?"},

example_text_completion.py

@@ -4,7 +4,7 @@
 import fire
 
 from llama import Llama
-
+from typing import List
 
 def main(
     ckpt_dir: str,
@@ -36,7 +36,7 @@ def main(
         max_batch_size=max_batch_size,
     )
 
-    prompts = [
+    prompts: List[str] = [
         # For these prompts, the expected answer is the natural continuation of the prompt
         "I believe the meaning of life is",
         "Simply put, the theory of relativity states that ",

llama/__init__.py

@@ -1,6 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # This software may be used and distributed according to the terms of the Llama 2 Community License Agreement.
 
-from .generation import Llama
+from .generation import Llama, Dialog
 from .model import ModelArgs, Transformer
 from .tokenizer import Tokenizer