Update app.py

app.py

@@ -559,51 +559,54 @@ def get_tensor_length(tensor):
         ret = ret * num
     return ret
 
-def summarize(
-    text, max_len=20, min_len=10
-):
+def _summarize(text, max_words=20):
+    prefix = "summarize: "
+    toks = tokenizer.encode( prefix + text, return_tensors="pt", truncation=False)
+    gen = model.generate(
+        toks,
+        length_penalty=2.0,
+        num_beams=max( get_tensor_length(toks) // 4 , 4 ),
+        early_stopping=True,
+        max_length=max_words
+    )
+    return tokenizer.decode(gen[0], skip_special_tokens=True)
+
+def _summ_step(length):
+    return max(length // 3 , min(10,length-4))
+
+def summarize(text, max_words=20):
     log(f'CALL summarize')
     
     words = text.split()
     
     if len(words) < 5:
-        print("Summarization Error: Text is too short, 5 words minimum!")
+        print("Summarization Error: Text is too short, 5 words minimum.")
         return text
         
-    prefix = "summarize: "
-    ret = ""
+    if max_words < 5 or max_words > 500:
+        print("Summarization Error: max_words value must be between 5 and 500 words.")
+        return text
 
-    for index in range(math.ceil( len(words) / 500 )):
-        
-        chunk = " ".join(words[ index*500:(index+1)*500 ])
-        inputs = tokenizer.encode( prefix + chunk, return_tensors="pt", truncation=False, add_special_tokens=True)
+    words_length = len(text.split())
+
+    if words_length >= 510:
+        shrink_step = 500 // (len(text.split()) / 500)
+        while words_length >= 510:
+            words = text.split()
+            above_limit = words[510:]
+            text = _summarize(
+                " ".join(words[0:510]), shrink_step
+            ) + " ".join(words[510:])
+            words_length = len(text.split())
+    
+    while words_length > max_words:
+        step = _summ_step(words_length)
+        mx = words_length - step
+        text = _summarize(text, mx)
+        words_length = len(text.split())
         
-        while get_tensor_length(inputs) > max_len:
-            
-            inputs = model.generate(
-                inputs,
-                length_penalty=2.0,
-                num_beams=4,
-                early_stopping=True,
-                max_length=max( get_tensor_length(inputs) // 4 , max_len ),
-                min_length=min_len
-            )
-
-        toks = tokenizer.decode(inputs[0], skip_special_tokens=True)
-        ret = ret + ("" if ret == "" else " ") + toks
-
-    inputs = tokenizer.encode( prefix + ret, return_tensors="pt", truncation=False)
-    gen = model.generate(
-        inputs,
-        length_penalty=1.0,
-        num_beams=4,
-        early_stopping=True,
-        max_length=max_len,
-        min_length=min_len
-    )
-    summary = tokenizer.decode(gen[0], skip_special_tokens=True)
-    log(f'RET summarize with summary as {summary}')
-    return summary
+    log(f'RET summarize with text as {text}')
+    return text
 
 def generate_random_string(length):
     characters = str(ascii_letters + digits)