rebase

llms/speculative_decoding/decoder.py

@@ -1,12 +1,12 @@
-import transformers
-from dataclasses import dataclass, field
-from model import Llama
-import mlx.core as mx
-import mlx.nn as nn
 import time
-import numpy as np
+from dataclasses import dataclass, field
 from typing import List, Optional
 
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+import transformers
+from model import Llama
 from prompts import create_urial_prompt
 
 

llms/speculative_decoding/main.py

@@ -1,4 +1,5 @@
 import time
+
 from decoder import SpeculativeDecoder
 
 # This will use the chat template from the primary model

llms/speculative_decoding/model.py

@@ -1,11 +1,10 @@
-from transformers import LlamaConfig, AutoModelForCausalLM
-import mlx.core as mx
-import mlx.nn as nn
-from mlx.utils import tree_unflatten, tree_map
-import mlx.core as mx
-import mlx.nn as nn
 from typing import Optional, Tuple
 
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_map, tree_unflatten
+from transformers import AutoModelForCausalLM, LlamaConfig
+
 
 def create_additive_causal_mask(N: int, offset: int = 0, dtype: mx.Dtype = mx.float32):
     rinds = mx.arange(offset + N)

llms/speculative_decoding/prompts.py

@@ -24,5 +24,6 @@ Sure, here are some of the most common types of renewable energy sources:
 Each type of renewable energy source has its own set of advantages and challenges, but collectively, they represent our best hope at achieving sustainable and environmentally friendly energy consumption. Please let me know if you have any other questions!
 ```"""
 
+
 def create_urial_prompt(message: str):
     return URIAL_1_SHOT + "\n\n# Query:\n```" + message + "```\n\n# Answer:\n```\n"