aboutsummaryrefslogtreecommitdiff
path: root/llama/__init__.py
diff options
context:
space:
mode:
Diffstat (limited to 'llama/__init__.py')
-rw-r--r--llama/__init__.py18
1 files changed, 18 insertions, 0 deletions
diff --git a/llama/__init__.py b/llama/__init__.py
new file mode 100644
index 0000000..2b9a9ef
--- /dev/null
+++ b/llama/__init__.py
@@ -0,0 +1,18 @@
+"""
+Llama2 model, loading infrastructure, and sampling helpers
+"""
+
+from .model import Llama
+
+from .generate import (
+ sample_sequence,
+ sample_batched_sequence,
+ generate_token_sequence,
+ generate_batched_token_sequence,
+ token_probabilities,
+ batched_token_probabilities
+)
+
+from .utils import (
+ load_llama_from_checkpoint
+)