make the use of adapter optional in generation

Author: remixer-dec

example.py

@@ -122,7 +122,7 @@ def main(
 
     with torch.inference_mode(mode=True):
         results = [generator.generate(
-            [prompt], max_gen_len=32, temperature=temperature, top_p=top_p
+            [prompt], max_gen_len=32, temperature=temperature, top_p=top_p, use_adapter=bool(adapter_path)
         ) for prompt in prompts]
 
     for result in results:

llama/generation.py

@@ -20,6 +20,7 @@ def generate(
         max_gen_len: int,
         temperature: float = 0.8,
         top_p: float = 0.95,
+        use_adapter: bool = True
     ) -> List[str]:
         bsz = len(prompts)
         params = self.model.params
@@ -39,7 +40,7 @@ def generate(
         start_pos = min_prompt_size
         prev_pos = 0
         for cur_pos in range(start_pos, total_len):
-            logits = self.model.forward(tokens[:, prev_pos:cur_pos], prev_pos)
+            logits = self.model.forward(tokens[:, prev_pos:cur_pos], prev_pos, use_adapter)
             if temperature > 0:
                 probs = torch.softmax(logits / temperature, dim=-1)
                 next_token = sample_top_p(probs, top_p)

llama/model.py

@@ -238,24 +238,25 @@ def __init__(self, params: ModelArgs):
         self.adapter_layer = params.adapter_layer
 
     @torch.inference_mode()
-    def forward(self, tokens: torch.Tensor, start_pos: int):
+    def forward(self, tokens: torch.Tensor, start_pos: int, use_adapter):
         _bsz, seqlen = tokens.shape
         h = self.tok_embeddings(tokens)
         #self.freqs_cis = self.freqs_cis.float().to(h.device)
         freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]
-        prompt = self.adapter_query.weight.reshape(self.params.adapter_layer, self.params.adapter_len, self.params.dim).unsqueeze(1)
 
         mask = None
         if seqlen > 1:
             mask = torch.full((1, 1, seqlen, seqlen), float("-inf"), device=torch.device('cpu'))
             mask = torch.triu(mask, diagonal=start_pos + 1).type_as(h)
 
-        for layer in self.layers[: -1 * self.params.adapter_layer]:
+        for layer in (self.layers[: -1 * self.params.adapter_layer]) if use_adapter else self.layers    :
             h = layer(h, start_pos, freqs_cis, (mask.to('mps') if mask is not None else None))
-        layer_index = 0
-        for layer in self.layers[-1 * self.params.adapter_layer:]:
-            h = layer(h, start_pos, freqs_cis, (mask.to('mps') if mask is not None else None), prompt[layer_index])
-            layer_index = layer_index + 1
+        if use_adapter:
+            prompt = self.adapter_query.weight.reshape(self.params.adapter_layer, self.params.adapter_len, self.params.dim).unsqueeze(1)
+            layer_index = 0
+            for layer in self.layers[-1 * self.params.adapter_layer:]:
+                h = layer(h, start_pos, freqs_cis, (mask.to('mps') if mask is not None else None), prompt[layer_index])
+                layer_index = layer_index + 1
         h = self.norm(h)
         output = self.output(h[:, -1, :])  # only compute last logits
         return output.float()