Update app.py

app.py

@@ -35,11 +35,12 @@ class GenerateRequest(BaseModel):
     stream: bool = True
     top_p: float = 1.0
     top_k: int = 50
-    repetition_penalty: float = 1.0
+    repetition_penalty: float = 1.1  # Increased default to discourage repetition
     num_return_sequences: int = 1
     do_sample: bool = True
     chunk_delay: float = 0.0
     stop_sequences: list[str] = []
+    no_repeat_ngram_size: int = 2 # Add parameter to prevent repeating ngrams
 
     @field_validator("model_name")
     def model_name_cannot_be_empty(cls, v):
@@ -112,6 +113,7 @@ async def generate(request: GenerateRequest):
         do_sample = request.do_sample
         chunk_delay = request.chunk_delay
         stop_sequences = request.stop_sequences
+        no_repeat_ngram_size = request.no_repeat_ngram_size
 
         model, tokenizer = await model_loader.load_model_and_tokenizer(model_name)
         device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -125,6 +127,7 @@ async def generate(request: GenerateRequest):
             repetition_penalty=repetition_penalty,
             do_sample=do_sample,
             num_return_sequences=num_return_sequences,
+            no_repeat_ngram_size=no_repeat_ngram_size, # Added to generation config
         )
 
         return StreamingResponse(
@@ -164,7 +167,8 @@ async def stream_text(model, tokenizer, input_text, generation_config, stop_sequ
                 stopping_criteria=stopping_criteria,
                 output_scores=True,
                 return_dict_in_generate=True,
-                pad_token_id=tokenizer.pad_token_id
+                pad_token_id=tokenizer.pad_token_id,
+                no_repeat_ngram_size=generation_config.no_repeat_ngram_size, # Passed to model.generate
             )
         except IndexError as e:
             print(f"IndexError during generation: {e}")