From 89092cdec82c3c1dbdb81cca97773b3829f59af8 Mon Sep 17 00:00:00 2001
From: tomorrmato <2ndorderode@gmail.com>
Date: Tue, 18 Jul 2023 23:14:58 -0700
Subject: [PATCH] added type hint in example code

---
 example_chat_completion.py | 6 +++---
 example_text_completion.py | 4 ++--
 llama/__init__.py          | 2 +-
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/example_chat_completion.py b/example_chat_completion.py
index 5043bc5..8f8de5e 100644
--- a/example_chat_completion.py
+++ b/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(
@@ -24,7 +24,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?"},
diff --git a/example_text_completion.py b/example_text_completion.py
index 4376b1e..228904d 100755
--- a/example_text_completion.py
+++ b/example_text_completion.py
@@ -4,7 +4,7 @@
 import fire
 
 from llama import Llama
-
+from typing import List
 
 def main(
     ckpt_dir: str,
@@ -22,7 +22,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 ",
diff --git a/llama/__init__.py b/llama/__init__.py
index 354342d..0bd1f86 100755
--- a/llama/__init__.py
+++ b/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