Update app.py

app.py

@@ -9,6 +9,7 @@ from tensorflow.keras.preprocessing.sequence import pad_sequences
 import time
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from transformers import GPT2Tokenizer, GPT2LMHeadModel
+import textwrap
 
 tok = GPT2Tokenizer.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2')
 model_checkpoint = 'cointegrated/rubert-tiny-toxicity'
@@ -52,17 +53,23 @@ def predict_text(text):
     return predicted_class
 
 
-def generate_text(model, prompt, max_length=150, temperature=1.0):
-    input_ids = tok.encode(prompt, return_tensors='pt')
-    output = model_finetuned.generate(
-        input_ids=input_ids,
-        max_length=max_length + len(input_ids[0]),
-        temperature=temperature,
-        num_return_sequences=1,
-        pad_token_id=tokenizer.eos_token_id
-    )
-    generated_text = textwrap.fill(tok.decode(out_), 60)
-    return generated_text
+def generate_text(model, prompt, max_length=150, temperature=1.0, num_beams=10, top_k=600, top_p=0.75, no_repeat_ngram_size=1, num_return_sequences=1):
+    input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
+    
+    with torch.inference_mode():
+        output = model.generate(
+            input_ids=input_ids,
+            max_length=max_length,
+            num_beams=num_beams,
+            do_sample=True,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            no_repeat_ngram_size=no_repeat_ngram_size,
+            num_return_sequences=num_return_sequences
+        )
+    texts = [textwrap.fill(tokenizer.decode(out), 60) for out in output]
+    return "\n------------------\n".join(texts)
 
 
 def page_reviews_classification():