Create handler.py

handler.py

@@ -0,0 +1,91 @@
+import torch
+from typing import Dict, List, Any
+import torch.nn as nn
+from transformers import GPT2LMHeadModel, GPT2Config, GPT2Tokenizer, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutput
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+# # get dtype
+# dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
+
+class CustomGPT2Model(PreTrainedModel):
+    def __init__(self, config):
+        super(CustomGPT2Model, self).__init__(config)
+
+        self.gpt2 = GPT2LMHeadModel.from_pretrained('gpt2-medium')
+        # Create an MLP layer to transform the ada-002 embedding to the GPT-2 hidden size
+        self.mlp = nn.Sequential(
+            nn.Linear(1536, 768),  # Adjust the hidden layer size as necessary
+            nn.ReLU(),
+            nn.Linear(768, config.n_embd)  # Adjust the output size to match GPT-2 embedding size
+        )
+
+        
+
+    def forward(self, inputs=None, ada_embedding=None, decoded_tkns=None, labels=None):
+        emb = self.mlp(ada_embedding)
+        emb = emb.unsqueeze(1)
+
+        if decoded_tkns is not None:
+            # Add the "encoded:" prefix, ada-002 embedding, "decoded:" prefix, and the decoded token
+            decoded_tkns = torch.cat([emb, self.gpt2.transformer.wte(decoded_tkns)], dim=1)
+        else:
+            decoded_tkns = emb
+
+        # Create the position ids
+        position_ids = torch.arange(0, decoded_tkns.size(1), dtype=torch.long).unsqueeze(0).to(emb.device)
+        # Forward the embeddings through the GPT-2 model with the correct position ids
+        outputs = self.gpt2(inputs_embeds=decoded_tkns, position_ids=position_ids)
+        logits = outputs.logits
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+        return CausalLMOutput(loss, logits, outputs.hidden_states)
+
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        # load the model
+        # Load the GPT-2 configuration
+        self.config = GPT2Config.from_pretrained('gpt2-medium')
+        
+        # Create the custom GPT-2 model and load the trained weights
+        self.model = CustomGPT2Model.from_pretrained(path, config=self.config)
+        
+        # Load the tokenizer
+        self.tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')
+
+
+    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
+        embedding = data.pop("embedding", None)
+        max_length=200
+        with torch.no_grad():
+            outputs = self.model(ada_embedding=embedding, decoded_tkns=None)
+        decoded_tkns = outputs.logits.argmax(dim=-1)
+    
+        for _ in range(max_length):
+            with torch.no_grad():
+                outputs = self.model(ada_embedding=embedding, decoded_tkns=decoded_tkns)
+            
+            # Get the most likely next token, sampled from top k
+            logits = outputs.logits[:, -1]
+            top_k_logits, top_k_indices = torch.topk(logits, k = 5)
+            next_token = torch.multinomial(F.softmax(top_k_logits, dim=-1), num_samples=1)
+            next_token = top_k_indices.gather(dim=1, index=next_token)
+        
+            if next_token[0].item() == self.tokenizer.eos_token_id:
+                break
+        
+            decoded_tkns = torch.cat((decoded_tkns, next_token), dim=1)
+    
+        # Convert the tensor of token IDs to a list of token IDs
+        token_ids = decoded_tkns[0].cpu().numpy().tolist()
+    
+        # Decode the token IDs back to a string
+        output_text = self.tokenizer.decode(token_ids, skip_special_tokens=True)
+        
+        return output_text