Update util.py

util.py

@@ -8,7 +8,6 @@ import torch
 import re
 
 
-
 def load_model(cp):
     tokenizer = AutoTokenizer.from_pretrained("VietAI/vit5-base")
     model  = AutoModelForSeq2SeqLM.from_pretrained(cp)
@@ -44,6 +43,8 @@ def get_subtitles(video_url):
     except Exception as e:
         return [], f"An error occurred: {e}"
 
+from youtube_transcript_api import YouTubeTranscriptApi
+
 
 def restore_punctuation(text):
     model = PunctuationModel()
@@ -115,12 +116,15 @@ def post_processing(text):
     text = " ".join(sentences)
     return text
 
+
 def display(text):
     sentences = re.split(r'(?<=[.!?])\s*', text)
     unique_sentences = list(dict.fromkeys(sentences[:-1]))
     formatted_sentences = [f"• {sentence}" for sentence in unique_sentences]
     return formatted_sentences
 
+
+
 def pipeline(url, model, tokenizer):
     trans, sub = get_subtitles(url)
     sub = restore_punctuation(sub)