Upload GPT124MTextGenerationPipeline

config.json

@@ -7,6 +7,21 @@
     "AutoModelForCausalLM": "modeling_gpt.GPTModelForTextGeneration"
   },
   "block_size": 1024,
+  "custom_pipelines": {
+    "text-generation": {
+      "default": {
+        "model": {
+          "pt": "samkeet/GPT_124M-Instruct"
+        }
+      },
+      "impl": "pipeline_gpt.GPT124MTextGenerationPipeline",
+      "pt": [
+        "AutoModelForCausalLM"
+      ],
+      "tf": [],
+      "type": "text"
+    }
+  },
   "model_type": "custom_gpt",
   "n_embd": 768,
   "n_head": 12,

pipeline_gpt.py

@@ -0,0 +1,59 @@
+# Importing necessary libraries
+import torch
+from transformers import Pipeline
+from .modeling_gpt import GPTModelForTextGeneration
+
+
+class GPT124MTextGenerationPipeline(Pipeline):
+
+    def _sanitize_parameters(self, **kwargs):
+        """
+        Organizes and sanitizes input parameters into separate dictionaries for:
+        - Preprocessing (encoding)
+        - Model forward pass (generation settings)
+        - Postprocessing (decoding)
+        """
+
+        preprocess_kwargs = {}
+        forward_kwargs = {
+            "max_length": kwargs.get("max_length", 50),
+            "do_sample": kwargs.get("do_sample", True),
+            "top_k": kwargs.get("top_k", 50),
+            "top_p": kwargs.get("top_p", 0.95),
+            "temperature": kwargs.get("temperature", 0.9),
+            "device": kwargs.get("device", self.device.type),
+        }
+        postprocess_kwargs = {}
+
+        return preprocess_kwargs, forward_kwargs, postprocess_kwargs
+
+    def preprocess(self, prompt_text: str, **preprocess_kwargs):
+        """
+        Encodes input text into token IDs using the tokenizer and converts it to a PyTorch tensor.
+        """
+
+        assert (
+            isinstance(prompt_text, str) and len(prompt_text) > 0
+        ), "prompt_text must be a non-empty string"
+
+        # Encode the input text
+        input_ids = self.tokenizer.encode(prompt_text)
+
+        # Convert to a PyTorch tensor
+        input_tensor = torch.tensor([input_ids])
+
+        return {"input_ids": input_tensor}
+
+    def _forward(self, model_inputs, **forward_kwargs):
+        """
+        Forwards the tokenized input to the model's generate method.
+        """
+
+        return self.model.generate(**model_inputs, **forward_kwargs)
+
+    def postprocess(self, model_output, **postprocess_kwargs):
+        """
+        Decodes token ID into human-readable text using the tokenizer.
+        """
+
+        return self.tokenizer.decode(model_output)