Update app.py

app.py

@@ -1,13 +1,26 @@
 import gradio as gr
-from transformers import pipeline
 import pytube
+from transformers import pipeline
+from time import sleep
+
 
 generator = pipeline('text-generation', model='parseny/youtube_comment_generation_02')
 
 def generate(text):
     yt = pytube.YouTube(text)
-    title = f"Title: {yt.title.lower()}\nComment:"
-    result = generator(title, max_new_tokens=30, num_return_sequences=1, do_sample=True)
+    while True:
+        try:
+            title = yt.title
+            break
+        except:
+            sleep(1)
+            yt = pytube.YouTube(text)
+            continue
+    title = title.lower()
+    title = title.replace("[^a-zA-Z#]", " ")
+    title = title.replace(" +", " ")
+    prompt = f"Title: {title}\nComment:"
+    result = generator(prompt, max_new_tokens=30, num_return_sequences=1, do_sample=True)
     return result[0]["generated_text"]
 
 examples = [