Update app.py

app.py

@@ -113,7 +113,7 @@ def get_base64_image(image):
     return f"data:image/jpeg;base64,{img_str}"
 
 @spaces.GPU(duration=120)
-def extract_vocabulary(pdf_text, images, translit_lang, translit_format, target_lang, max_text_char=1500):
+def extract_vocabulary(pdf_text, images, translit_lang, translit_format, target_lang, max_text_char=1500, repetition_penalty_val=1.1):
     """Use Transformers to extract vocabulary from text and images."""
     global model, processor
     
@@ -187,7 +187,7 @@ Text:
             streamer=streamer,
             max_new_tokens=2048*16,
             do_sample=True,
-            repetition_penalty=1.1,
+            repetition_penalty=repetition_penalty_val,
         )
 
         if len(images) > 0:
@@ -231,7 +231,7 @@ Text:
         print(f"Error parsing JSON: {e}\nRaw output: {output_text}")
         yield output_text, []
 
-def translate_vocabulary(korean_words, translit_lang, translit_format, target_lang):
+def translate_vocabulary(korean_words, translit_lang, translit_format, target_lang, repetition_penalty_val=1.1):
     """Use Transformers text-only inference to translate/transliterate Korean words."""
     global model, processor
     
@@ -282,7 +282,7 @@ Korean words:
             # top_p=0.95,
             temperature=1.0, top_p=0.95, top_k=20, min_p=0.0, 
             # presence_penalty=1.5, 
-            repetition_penalty=1.1,
+            repetition_penalty=repetition_penalty_val,
             do_sample=True
         )
         
@@ -331,7 +331,7 @@ def hash_file(filepath):
         return hashlib.md5(f.read(1024*1024)).hexdigest()
 
 @spaces.GPU(duration=120)
-def process_pdf(pdf_file, url_input, translit_lang, translit_format, target_lang, max_text_char, last_source_hash, last_korean_words, progress=gr.Progress()):
+def process_pdf(pdf_file, url_input, translit_lang, translit_format, target_lang, max_text_char, repetition_penalty_val, last_source_hash, last_korean_words, progress=gr.Progress()):
     global tts, voice_style
     
     # Clean language choices from "Family - Language" to just "Language"
@@ -359,7 +359,7 @@ def process_pdf(pdf_file, url_input, translit_lang, translit_format, target_lang
     #     progress(0.2, desc="Translating previously extracted vocabulary...")
     #     korean_words = [item.get("korean") for item in last_korean_words if item.get("korean")]
     #     for attempt in range(1, 4):
-    #         vocab_list = translate_vocabulary(korean_words, translit_lang, translit_format, target_lang)
+    #         vocab_list = translate_vocabulary(korean_words, translit_lang, translit_format, target_lang, repetition_penalty_val)
     #         if vocab_list:
     #             break
     # else:
@@ -383,7 +383,7 @@ def process_pdf(pdf_file, url_input, translit_lang, translit_format, target_lang
     stream_text = ""
     for attempt in range(1, 4):
         progress(0.2, desc=f"Extracting vocabulary (Attempt {attempt}/3)...")
-        for stream_t, v_list in extract_vocabulary(content_text, images, translit_lang, translit_format, target_lang, max_text_char):
+        for stream_t, v_list in extract_vocabulary(content_text, images, translit_lang, translit_format, target_lang, max_text_char, repetition_penalty_val):
             stream_text = stream_t
             if v_list is not None:
                 vocab_list = v_list
@@ -758,6 +758,7 @@ def create_demo():
                     value="Indo-European - English"
                 )
                 max_text_char_input = gr.Slider(minimum=1000, maximum=30000, step=1000, value=1500, label="Max Input Text Length (Characters)")
+                repetition_penalty_input = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.2, label="Repetition Penalty")
                 
                 with gr.Row():
                     submit_btn = gr.Button("✨ Generate Flashcards ✨", variant="primary")
@@ -776,7 +777,7 @@ def create_demo():
 
         generate_event = submit_btn.click(
             fn=process_pdf, 
-            inputs=[pdf_input, url_input, translit_lang, translit_format, target_lang, max_text_char_input, last_source_state, last_korean_words_state], 
+            inputs=[pdf_input, url_input, translit_lang, translit_format, target_lang, max_text_char_input, repetition_penalty_input, last_source_state, last_korean_words_state], 
             outputs=[output_html, last_source_state, last_korean_words_state, stream_box, extracted_text_box, extracted_images_gallery]
         )