NeuraCraft commited on
Commit
2334367
·
1 Parent(s): ce83413

Upload lance_ai_model.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. lance_ai_model.py +13 -1
lance_ai_model.py CHANGED
@@ -348,8 +348,20 @@ class LanceAI(LanceAIPreTrainedModel, GenerationMixin):
348
 
349
  return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=past)
350
 
 
 
 
 
 
 
 
 
 
 
 
 
351
  def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
352
- if past_key_values is not None:
353
  input_ids = input_ids[:, -1:]
354
  return {
355
  "input_ids": input_ids,
 
348
 
349
  return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=past)
350
 
351
+ def _past_seq_len(self, past_key_values):
352
+ if past_key_values is None:
353
+ return 0
354
+ if hasattr(past_key_values, 'get_seq_length'):
355
+ return past_key_values.get_seq_length()
356
+ if isinstance(past_key_values, list) and len(past_key_values) > 0:
357
+ if past_key_values[0] is not None:
358
+ k, v = past_key_values[0]
359
+ if k is not None:
360
+ return k.shape[2]
361
+ return 0
362
+
363
  def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
364
+ if self._past_seq_len(past_key_values) > 0:
365
  input_ids = input_ids[:, -1:]
366
  return {
367
  "input_ids": input_ids,