Update gen.py

gen.py

@@ -7,7 +7,9 @@ from pathlib import Path
 
 import torch
 from fairscale.nn.model_parallel.initialize import initialize_model_parallel
-from llama import ModelArgs, Transformer, Tokenizer, LLaMA
+from llama.generation import LLaMA
+from llama.model import ModelArgs, Transformer
+from llama.tokenizer import Tokenizer
 
 from google.cloud import storage