Upload folder using huggingface_hub

README.md

@@ -1,9 +1,3 @@
----
-license: apache-2.0
-datasets:
-- tuetschek/atis
-language:
-- en
-pipeline_tag: text-generation
-library_name: transformers
----
+---
+license: apache-2.0
+---

config.json

@@ -0,0 +1,19 @@
+{
+  "architectures": [
+    "TransformerChatbot"
+  ],
+  "d_ff": 2048,
+  "d_model": 512,
+  "dropout": 0.1,
+  "max_len": 5000,
+  "max_turns": 16,
+  "model_type": "transformer-chat",
+  "num_decoder_layers": 6,
+  "num_encoder_layers": 6,
+  "num_heads": 8,
+  "num_roles": 2,
+  "num_slots": 22,
+  "torch_dtype": "float32",
+  "transformers_version": "4.52.4",
+  "vocab_size": 25000
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64652386a8c09a9a7da5d5f5490ef5c9c32ab8025133373bd37fe22061ab9f3a
+size 412838448

save_hf_model.py

@@ -0,0 +1,30 @@
+from transformer_chat_config import TransformerChatConfig
+from transformer_chat import TransformerChatbot
+import torch
+import os
+
+config = TransformerChatConfig(
+    vocab_size=25000,  
+    d_model=512,
+    num_heads=8,
+    d_ff=2048,
+    num_encoder_layers=6,
+    num_decoder_layers=6,
+    num_roles=2,
+    max_turns=16,
+    num_slots=22,
+    dropout=0.1,
+    max_len=5000
+)
+
+model = TransformerChatbot(config)
+
+# Load trained weights
+model.load_state_dict(torch.load("atis_transformer.pt", map_location="cpu"))
+
+# Save model and config in HuggingFace format
+save_dir = os.path.join("..", "my-hf-chatbot")
+model.save_pretrained(save_dir)
+config.save_pretrained(save_dir)
+
+print(f"Model and config saved to {save_dir}") 

special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

tokenization/save_hf_tokenizer.py

@@ -0,0 +1,29 @@
+from transformers import PreTrainedTokenizerFast, AutoTokenizer
+
+# Load the tokenizer.json
+hf_tokenizer = PreTrainedTokenizerFast(tokenizer_file="tokenizer.json")
+
+# Set special tokens
+hf_tokenizer.pad_token = "[PAD]"
+hf_tokenizer.unk_token = "[UNK]"
+hf_tokenizer.cls_token = "[CLS]"
+hf_tokenizer.sep_token = "[SEP]"
+hf_tokenizer.mask_token = "[MASK]"
+
+# Save in HuggingFace format
+hf_tokenizer.save_pretrained("../my-hf-chatbot")
+
+# Test loading
+loaded_tokenizer = AutoTokenizer.from_pretrained("../my-hf-chatbot")
+print("Loaded tokenizer vocab size:", len(loaded_tokenizer))
+print("Special tokens:")
+print("PAD:", loaded_tokenizer.pad_token, loaded_tokenizer.pad_token_id)
+print("UNK:", loaded_tokenizer.unk_token, loaded_tokenizer.unk_token_id)
+print("CLS:", loaded_tokenizer.cls_token, loaded_tokenizer.cls_token_id)
+print("SEP:", loaded_tokenizer.sep_token, loaded_tokenizer.sep_token_id)
+print("MASK:", loaded_tokenizer.mask_token, loaded_tokenizer.mask_token_id)
+
+# Test encoding
+sample = "Let's test this tokenizer."
+enc = loaded_tokenizer(sample)
+print("Encoded tokens:", enc.tokens()) 

tokenizer_config.json

@@ -0,0 +1,60 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "max_length": 128,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_to_multiple_of": null,
+  "pad_token": "[PAD]",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "sep_token": "[SEP]",
+  "stride": 0,
+  "tokenizer_class": "PreTrainedTokenizer",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "[UNK]"
+}

transformer_chat.py

@@ -1,6 +1,8 @@
 import torch
 import torch.nn as nn
 import math
+from transformers import PreTrainedModel
+from transformer_chat_config import TransformerChatConfig
 
 # Import neccessary layers
 from built_transformer.embeddings import Embeddings
@@ -9,69 +11,53 @@ from built_transformer.decoders import Decoder, DecoderLayer
 from built_transformer.positional_encodings import PositionalEncoding
 from built_transformer.slot_classifier import SlotClassifier
 
-class TransformerChatbot(nn.Module):
+class TransformerChatbot(PreTrainedModel):
     """
     Unified Transformer-based chatbot model that combines:
     - Joint token/role/turn embeddings
     - Encoder-decoder architecture with attention
     - Slot-filling classification
-    - Generation capabilities
+    - Generation capabilities (newly added, hopefully compatible with HuggingFace)
     """
-    def __init__(
-        self,
-        vocab_size: int,
-        d_model: int = 512,
-        num_heads: int = 8,
-        d_ff: int = 2048,
-        num_encoder_layers: int = 6,
-        num_decoder_layers: int = 6,
-        num_roles: int = 2,
-        max_turns: int = 16,
-        num_slots: int = 4,
-        dropout: float = 0.1,
-        max_len: int = 5000
-    ):
-        super().__init__()
-        
+    config_class = TransformerChatConfig
+
+    def __init__(self, config: TransformerChatConfig):
+        super().__init__(config)
+        self.config = config
         # Embeddings for tokens, roles, and turns
         self.embed = Embeddings(
-            char=vocab_size, # Fixed type and name mismatch
-            dimension_for_model=d_model,
-            num_of_roles=num_roles,
-            max_turns=max_turns
+            char=config.vocab_size,
+            dimension_for_model=config.d_model,
+            num_of_roles=config.num_roles,
+            max_turns=config.max_turns
         )
-        
         # Positional encoding
-        self.pos_enc = PositionalEncoding(d_model, dropout, max_len)
-        
+        self.pos_enc = PositionalEncoding(config.d_model, config.dropout, config.max_len)
         # Encoder stack
         self.encoder = Encoder(
-            vocab_size=vocab_size,
-            dimension_of_model=d_model,
-            num_of_heads=num_heads,
-            num_layers=num_encoder_layers,
-            dim_feedforward=d_ff,
-            dropout=dropout,
-            max_len=max_len,
-            num_of_roles=num_roles,
-            max_turns=max_turns
+            vocab_size=config.vocab_size,
+            dimension_of_model=config.d_model,
+            num_of_heads=config.num_heads,
+            num_layers=config.num_encoder_layers,
+            dim_feedforward=config.d_ff,
+            dropout=config.dropout,
+            max_len=config.max_len,
+            num_of_roles=config.num_roles,
+            max_turns=config.max_turns
         )
-        
         # Decoder stack
         self.decoder = Decoder(
-            vocab_size=vocab_size,
-            dimension_for_model=d_model,
-            num_layers=num_decoder_layers,
-            num_of_heads=num_heads,
-            dim_feedforward=d_ff,
-            dropout=dropout,
-            max_len=max_len
+            vocab_size=config.vocab_size,
+            dimension_for_model=config.d_model,
+            num_layers=config.num_decoder_layers,
+            num_of_heads=config.num_heads,
+            dim_feedforward=config.d_ff,
+            dropout=config.dropout,
+            max_len=config.max_len
         )
-        
         # Output projections
-        self.out_proj = nn.Linear(d_model, vocab_size)
-        self.slot_classifier = SlotClassifier(d_model, num_slots)
-        
+        self.out_proj = nn.Linear(config.d_model, config.vocab_size)
+        self.slot_classifier = SlotClassifier(config.d_model, config.num_slots)
         # Initialize parameters
         self._init_parameters()
         
@@ -90,8 +76,8 @@ class TransformerChatbot(nn.Module):
             # Copy the old embedding weights to the new structure
             state_dict['encoder.embed.lut.weight'] = old_embed_weight
             # Initialize role and turn embeddings with correct sizes
-            state_dict['encoder.embed.lut_roles.weight'] = torch.zeros(2, old_embed_weight.size(1))  # 2 roles
-            state_dict['encoder.embed.lut_turns.weight'] = torch.zeros(16, old_embed_weight.size(1))  # 16 turns
+            state_dict['encoder.embed.lut_roles.weight'] = torch.zeros(self.config.num_roles, old_embed_weight.size(1))
+            state_dict['encoder.embed.lut_turns.weight'] = torch.zeros(self.config.max_turns, old_embed_weight.size(1))
             state_dict['encoder.embed.norm.weight'] = torch.ones(old_embed_weight.size(1))
             state_dict['encoder.embed.norm.bias'] = torch.zeros(old_embed_weight.size(1))
             
@@ -145,34 +131,51 @@ class TransformerChatbot(nn.Module):
     
     def forward(
         self,
-        src_tokens,
-        tgt_tokens,
-        src_roles,
-        tgt_roles,
-        src_turns,
-        tgt_turns,
-        src_mask=None,
-        tgt_mask=None
+        input_ids=None,
+        decoder_input_ids=None,
+        attention_mask=None,
+        decoder_attention_mask=None,
+        src_roles=None,
+        tgt_roles=None,
+        src_turns=None,
+        tgt_turns=None,
+        **kwargs
     ):
         """
-        Full forward pass combining encoding, decoding, and slot classification.
+        HuggingFace-compatible forward method.
         Args:
-            src_tokens: [B, S] source token IDs
-            tgt_tokens: [B, T] target token IDs
-            src_roles:  [B, S] source role IDs
-            tgt_roles:  [B, T] target role IDs
-            src_turns:  [B, S] source turn IDs
-            tgt_turns:  [B, T] target turn IDs
-            src_mask:   [B, 1, 1, S] source mask
-            tgt_mask:   [B, 1, T, T] target mask
+            input_ids: [B, S] source token IDs
+            decoder_input_ids: [B, T] target token IDs
+            attention_mask: [B, 1, 1, S] source mask
+            decoder_attention_mask: [B, 1, T, T] target mask
+            src_roles: [B, S] source role IDs
+            tgt_roles: [B, T] target role IDs
+            src_turns: [B, S] source turn IDs
+            tgt_turns: [B, T] target turn IDs
         Returns:
             gen_logits: [B, T, vocab_size] generation logits
             slot_logits: [B, num_slots] slot classification logits
         """
-        # Encode source sequence
+        src_tokens = input_ids
+        tgt_tokens = decoder_input_ids
+        src_mask = attention_mask
+        tgt_mask = decoder_attention_mask
+
+        # Infer roles and turns if not provided
+        if src_tokens is not None:
+            batch_size, src_len = src_tokens.shape
+            if src_roles is None:
+                src_roles = torch.zeros_like(src_tokens)
+            if src_turns is None:
+                src_turns = torch.arange(src_len, device=src_tokens.device).unsqueeze(0).expand(batch_size, src_len)
+        if tgt_tokens is not None:
+            batch_size, tgt_len = tgt_tokens.shape
+            if tgt_roles is None:
+                tgt_roles = torch.zeros_like(tgt_tokens)
+            if tgt_turns is None:
+                tgt_turns = torch.arange(tgt_len, device=tgt_tokens.device).unsqueeze(0).expand(batch_size, tgt_len)
+
         enc_out = self.encode(src_tokens, src_roles, src_turns, src_mask)
-        
-        # Decode target sequence
         gen_logits = self.decode(
             tgt_tokens,
             enc_out,
@@ -181,9 +184,6 @@ class TransformerChatbot(nn.Module):
             src_mask,
             tgt_mask
         )
-        
-        # Use first position of encoder output for slot classification
         cls_rep = enc_out[:, 0, :]
         slot_logits = self.slot_classifier(cls_rep)
-        
-        return gen_logits, slot_logits
+        return {"gen_logits": gen_logits, "slot_logits": slot_logits}

transformer_chat_config.py

@@ -0,0 +1,31 @@
+from transformers import PretrainedConfig
+
+class TransformerChatConfig(PretrainedConfig):
+    model_type = "transformer-chat"
+    def __init__(
+        self,
+        vocab_size=30522,
+        d_model=512,
+        num_heads=8,
+        d_ff=2048,
+        num_encoder_layers=6,
+        num_decoder_layers=6,
+        num_roles=2,
+        max_turns=16,
+        num_slots=4,
+        dropout=0.1,
+        max_len=5000,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.d_ff = d_ff
+        self.num_encoder_layers = num_encoder_layers
+        self.num_decoder_layers = num_decoder_layers
+        self.num_roles = num_roles
+        self.max_turns = max_turns
+        self.num_slots = num_slots
+        self.dropout = dropout
+        self.max_len = max_len