Update main.py

main.py

@@ -15,6 +15,7 @@ class Item(BaseModel):
     temperature: float = 0.6
     max_new_tokens: int = 1024
     top_p: float = 0.95
+    repetition_penalty: float = 1.0
     seed : int = 42
     
 app = FastAPI()
@@ -56,7 +57,7 @@ def generate(item: Item):
         tokenizer.convert_tokens_to_ids("<|eot_id|>")
     ]
 
-    outputs = model.generate(input_ids, eos_token_id=terminators, do_sample=True)
+    outputs = model.generate(input_ids, eos_token_id=terminators, do_sample=True,  **generate_kwargs,)
     response = outputs[0][input_ids.shape[-1]:]
     return tokenizer.decode(response, skip_special_tokens=True)