update

app.py

@@ -17,6 +17,7 @@ def infer(
     top_p=1.0,
     top_k=40,
     num_completions=1,
+    repetition_penalty=1.0,
     seed=42,
     stop="\n"
 ):
@@ -28,6 +29,7 @@ def infer(
     temperature = float(temperature)
     top_p = float(top_p)
     top_k = int(top_k)
+    repetition_penalty = float(repetition_penalty)
     stop = stop.split(";")
     seed = seed
     
@@ -36,6 +38,7 @@ def infer(
     assert 0.0 <= temperature <= 10.0
     assert 0.0 <= top_p <= 1.0
     assert 1 <= top_k <= 1000
+    assert 0.9 <= repetition_penalty <= 3.0
 
     if temperature == 0.0:
         temperature = 0.01
@@ -48,6 +51,7 @@ def infer(
         "top_k": top_k,
         "temperature": temperature,
         "max_tokens": max_new_tokens,
+        "repetition_penalty": repetition_penalty,
         "stop": stop,
     }
     print(f"send: {datetime.now()}")
@@ -223,6 +227,7 @@ def main():
     if 'preset' not in st.session_state:
         st.session_state.preset = "Sentiment Analysis"
         st.session_state.top_k = "40"
+        st.session_state.repetition_penalty = "1.0"
         st.session_state.stop = r'\n'
         set_preset()
 
@@ -252,6 +257,7 @@ def main():
         top_p = st.text_input('top_p', st.session_state.top_p)
         # num_completions = st.text_input('num_completions (only the best one will be returend)', "1")
         num_completions = "1"
+        repetition_penalty = st.text_input('repetition_penalty', st.session_state.repetition_penalty)
         stop = st.text_input('stop, split by;', st.session_state.stop)
         # seed = st.text_input('seed', "42")
         seed = "42"
@@ -275,7 +281,8 @@ def main():
             generated_area.markdown("<b>" + to_md(prompt) + "</b>", unsafe_allow_html=True)
             report_text = infer(
                 prompt, model_name=model_name, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k,
-                num_completions=num_completions, seed=seed, stop=literal_eval("'''"+stop+"'''"),
+                num_completions=num_completions, repetition_penalty=repetition_penalty,
+                seed=seed, stop=literal_eval("'''"+stop+"'''"),
             )
             generated_area.markdown("<b>" + to_md(prompt) + "</b><mark style='background-color: #cbeacd'>" + to_md(report_text)+"</mark>", unsafe_allow_html=True)