Create inference.py

code/inference.py

@@ -0,0 +1,37 @@
+import torch
+import torch.nn.functional as F
+
+
+def generate(model,
+             tokenizer,
+             prompt: str,
+             n_tokens_to_gen: int = 200,
+             sample: bool = True,
+             top_k: int = 40):
+    model.eval()
+
+    input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to("cuda")
+
+    for token_n in range(n_tokens_to_gen):
+        with torch.no_grad():
+            indices_to_input = input_ids
+            next_token_logits = mamba_model(indices_to_input)[:, -1]
+
+        probs = F.softmax(next_token_logits, dim=-1)
+        (batch, vocab_size) = probs.shape
+
+        if top_k is not None:
+            (values, indices) = torch.topk(probs, k=top_k)
+            probs[probs < values[:, -1, None]] = 0
+            probs = probs / probs.sum(axis=1, keepdims=True)
+
+        if sample:
+            next_indices = torch.multinomial(probs, num_samples=1)
+        else:
+            next_indices = torch.argmax(probs, dim=-1)[:, None]
+
+        input_ids = torch.cat([input_ids, next_indices], dim=1)
+
+    output_completions = [tokenizer.decode(output.tolist()) for output in input_ids][0]
+
+    return output_completions