Upload LanceAI

generation_config.json

@@ -1,4 +1,11 @@
 {
   "_from_model_config": true,
+  "do_sample": true,
+  "max_new_tokens": 150,
+  "no_repeat_ngram_size": 3,
+  "repetition_penalty": 1.3,
+  "temperature": 0.8,
+  "top_k": 40,
+  "top_p": 0.9,
   "transformers_version": "4.51.3"
 }

lance_ai_model.py

@@ -34,9 +34,19 @@ class LanceAI(PreTrainedModel, GenerationMixin):
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)
         self.loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
 
+        # Set generation config defaults for more natural responses
+        self.generation_config.max_new_tokens = 150
+        self.generation_config.temperature = 0.8
+        self.generation_config.top_k = 40
+        self.generation_config.top_p = 0.9
+        self.generation_config.do_sample = True
+        self.generation_config.repetition_penalty = 1.3
+        self.generation_config.no_repeat_ngram_size = 3
+        self.generation_config.length_penalty = 1.0
+
         self.init_weights()
 
-    def forward(self, input_ids=None, attention_mask=None, labels=None, inputs_embeds=None, return_dict=True, use_cache=False):
+    def forward(self, input_ids=None, attention_mask=None, labels=None, inputs_embeds=None, return_dict=True, use_cache=False, **kwargs):
         embeddings = self.embedding(input_ids) if inputs_embeds is None else inputs_embeds
         encoder_output = self.encoder(embeddings)
         decoder_output = self.decoder(embeddings, encoder_output)
@@ -55,6 +65,25 @@ class LanceAI(PreTrainedModel, GenerationMixin):
 
         return (loss, logits) if loss is not None else logits
 
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
+        # Only last token for inputs_ids if past is defined in kwargs
+        if past_key_values:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "past_key_values": past_key_values,
+            **kwargs,
+        }
+    
+    def _reorder_cache(self, past_key_values, beam_idx):
+        # Reorder the cache for beam search
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
+        return reordered_past
+
 # Register with Hugging Face
 CONFIG_MAPPING.register("lance_ai", LanceAIConfig)
 MODEL_FOR_CAUSAL_LM_MAPPING.register(LanceAIConfig, LanceAI)