repetition penalty

config.json

@@ -1,11 +1,11 @@
 {
-  "_name_or_path": "vikhyatk/moondream2",
+  "_name_or_path": "gokaygokay/moondream-prompt",
   "architectures": [
     "Moondream"
   ],
   "auto_map": {
-    "AutoConfig": "vikhyatk/moondream2--configuration_moondream.MoondreamConfig",
-    "AutoModelForCausalLM": "vikhyatk/moondream2--moondream.Moondream"
+    "AutoConfig": "gokaygokay/moondream-prompt--configuration_moondream.MoondreamConfig",
+    "AutoModelForCausalLM": "gokaygokay/moondream-prompt--moondream.Moondream"
   },
   "model_type": "moondream1",
   "phi_config": {

moondream.py

@@ -64,6 +64,7 @@ class Moondream(PreTrainedModel):
         tokenizer,
         eos_text="<END>",
         max_new_tokens=128,
+        repetition_penalty=1.5,
         **kwargs,
     ):
         eos_tokens = tokenizer(eos_text, add_special_tokens=False)[0].ids
@@ -73,6 +74,7 @@ class Moondream(PreTrainedModel):
             "bos_token_id": tokenizer.bos_token_id,
             "pad_token_id": tokenizer.eos_token_id,
             "max_new_tokens": max_new_tokens,
+            "repetition_penalty": repetition_penalty,
             **kwargs,
         }
 
@@ -90,6 +92,8 @@ class Moondream(PreTrainedModel):
         question,
         tokenizer,
         chat_history="",
+        repetition_penalty=1.5,
+        max_new_tokens=128,
         result_queue=None,
         **kwargs,
     ):
@@ -99,7 +103,8 @@ class Moondream(PreTrainedModel):
             prompt,
             eos_text="<END>",
             tokenizer=tokenizer,
-            max_new_tokens=512,
+            max_new_tokens=max_new_tokens,
+            repetition_penalty=repetition_penalty,
             **kwargs,
         )[0]
         cleaned_answer = re.sub("<$|<END$", "", answer).strip()
@@ -115,6 +120,8 @@ class Moondream(PreTrainedModel):
         images,
         prompts,
         tokenizer,
+        repetition_penalty=1.5,
+        max_new_tokens=128,
         **kwargs,
     ):
         eos_tokens = tokenizer("<END>", add_special_tokens=False)[0].ids
@@ -162,7 +169,8 @@ class Moondream(PreTrainedModel):
             "eos_token_id": eos_tokens,
             "bos_token_id": tokenizer.bos_token_id,
             "pad_token_id": tokenizer.eos_token_id,
-            "max_new_tokens": 512,
+            "max_new_tokens": max_new_tokens,
+            "repetition_penalty": repetition_penalty,
             **kwargs,
         }