Update model_class.py

model_class.py

@@ -174,3 +174,77 @@ class CustomGPT(
         x = self.final_norm(x)
         logits = self.out_head(x) #[2,4,50257]
         return logits
+
+    def format_input(self, entry):
+        instruction_text = (
+            f"Below is an instruction that describes a task. "
+            f"Write a response that appropriately completes the request."
+            f"\n\n### Instruction:\n{entry['instruction']}"
+        )
+        input_text = f"\n\n### Input:\n{entry['input']}" if entry["input"] else ""
+        return instruction_text + input_text
+
+    def text_to_token_ids(self, text, tokenizer):
+        encoded = tokenizer.encode(text, allowed_special={'<|endoftext|>'})
+        encoded_tensor = torch.tensor(encoded).unsqueeze(0) # add batch dimension
+        return encoded_tensor
+
+    def token_ids_to_text(self, token_ids, tokenizer):
+        flat = token_ids.squeeze(0) # remove batch dimension
+        return tokenizer.decode(flat.tolist())
+
+    def generate(self, idx, max_new_tokens, context_size, temperature=0.0, top_k=None, eos_id=None):
+        for _ in range(max_new_tokens):
+            idx_cond = idx[:, -context_size:]
+            with torch.no_grad():
+                logits = self(idx_cond)
+            logits = logits[:, -1, :]
+    
+            if top_k is not None:
+                # Keep only top_k values
+                top_logits, _ = torch.topk(logits, top_k)
+                min_val = top_logits[:, -1] # select the last element i.e., the smallest from each batch's output
+                logits = torch.where(logits < min_val, torch.tensor(float("-inf")).to(logits.device), logits)
+    
+            # New: Apply temperature scaling
+            if temperature > 0.0:
+                logits = logits / temperature
+    
+                # Apply softmax to get probabilities
+                probs = torch.softmax(logits, dim=-1)  # (batch_size, context_len)
+    
+                # Sample from the distribution
+                idx_next = torch.multinomial(probs, num_samples=1)  # (batch_size, 1)
+    
+            # Otherwise same as before: get idx of the vocab entry with the highest logits value
+            else:
+                idx_next = torch.argmax(logits, dim=-1, keepdim=True)  # (batch_size, 1)
+    
+            if idx_next == eos_id:  # Stop generating early if end-of-sequence token is encountered and eos_id is specified
+                break
+    
+            # Same as before: append sampled index to the running sequence
+            idx = torch.cat((idx, idx_next), dim=1)  # (batch_size, num_tokens+1)
+    
+        return idx
+
+    def generate_response(self, input_entry, temperature=0.0, topk=None):
+        current_device = next(self.parameters()).device
+        self.eval()
+        input_text = self.format_input(entry)
+
+        token_ids = generate(
+            idx=self.text_to_token_ids(input_text, tokenizer).to(current_device),
+            max_new_tokens=256,
+            context_size=1024,
+            temperatue=temperature,
+            topk=topk,
+            eos_id=50256
+        )
+        generated_text = self.token_ids_to_text(token_ids, tokenizer)
+        response_text = (
+            generated_text[len(input_text):]
+            .replace("### Response:", "")
+            .strip()
+            )
+        return response_text.strip()