Update model.safetensors

model.safetensors

@@ -1,22 +1,34 @@
 import torch
-from transformers import PreTrainedModel, PreTrainedTokenizerFast, PretrainedConfig
-from transformers.modeling_outputs import ModelOutput
+from transformers import PreTrainedModel, PreTrainedTokenizerFast, PretrainedConfig, LogitsProcessorList
+from transformers.generation_utils import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutput
 
 class HelloWorldConfig(PretrainedConfig):
     model_type = "hello-world"
 
-class HelloWorldModel(PreTrainedModel):
+class HelloWorldModel(PreTrainedModel, GenerationMixin):
     config_class = HelloWorldConfig
 
     def __init__(self, config):
         super().__init__(config)
 
-    def forward(self, *args, **kwargs):
-        return ModelOutput(logits=torch.tensor([[0]]), decoder_hidden_states=["Hello, world!"])
+    def forward(self, input_ids=None, **kwargs):
+        batch_size = input_ids.shape[0]
+        sequence_length = input_ids.shape[1]
+        
+        # Generate a tensor with repeated "Hello, world!" token IDs
+        hello_world_token_id = self.config.vocab_size - 1  # assuming last token is "Hello, world!"
+        logits = torch.full((batch_size, sequence_length, self.config.vocab_size), float('-inf'))
+        logits[:, :, hello_world_token_id] = 0  # setting logits for "Hello, world!" to 0 (highest value)
+        
+        return CausalLMOutput(logits=logits)
 
-tokenizer = PreTrainedTokenizerFast(tokenizer_file="tokenizer.json")
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+    def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder=False):
+        return model_kwargs
 
-# Dummy tokenizer configuration to work with the model
 tokenizer_config = {
     "do_lower_case": False,
     "model_max_length": 512,
@@ -25,9 +37,14 @@ tokenizer_config = {
     "tokenizer_file": "tokenizer.json",
     "unk_token": "<unk>",
     "bos_token": "<s>",
-    "eos_token": "</s>"
+    "eos_token": "</s>",
+    "vocab_size": 1,  # Simplified vocabulary size
 }
 
+# Save tokenizer configuration
 with open("tokenizer.json", "w") as f:
     import json
     json.dump(tokenizer_config, f)
+
+tokenizer = PreTrainedTokenizerFast(tokenizer_file="tokenizer.json")
+tokenizer.add_tokens(["Hello, world!"])