Add punctuation part

app.py

@@ -1,31 +1,13 @@
+from os import pipe
 import gradio as gr
-import requests
-from youtube_transcript_api import YouTubeTranscriptApi
-import json
+from functions.punctuation import punctuate
 
 
-def retrieve_url(vid_id):
-    try:
-        transcript = YouTubeTranscriptApi.get_transcript(vid_id)
-    except Exception as e:
-        raise e
-    req = requests.get(
-        f"https://yt.lemnoslife.com/noKey/videos?part=snippet&id={vid_id}"
-    )
-    if req.status_code == 200:
-        information = json.loads(req.content)
-    else:
-        #        print(req.status_code)
-        information = None
-        return ""
-    #    print(transcript)
-    text = " ".join([x["text"] for x in transcript])
-    return text
+def pipeline(video_id):
+    punctuated_text = punctuate(video_id)
+    return punctuated_text
 
 
-def greet(name):
-    return "Hello " + name + "!!"
-
-
-demo = gr.Interface(fn=retrieve_url, inputs="text", outputs="text")
+# print(pipeline("VL5M5ZihJK4"))
+demo = gr.Interface(fn=pipeline, inputs="text", outputs="text")
 demo.launch(share=True)

functions/punctuation.py

@@ -0,0 +1,58 @@
+import requests
+from youtube_transcript_api import YouTubeTranscriptApi
+import json
+import os
+
+headers = {
+    "Authorization": f"Bearer {os.environ['HF_Token']}"
+}  # NOTE: put this somewhere else
+
+
+def retrieve_transcript(vid_id):
+    try:
+        transcript = YouTubeTranscriptApi.get_transcript(vid_id)
+        return transcript
+    except Exception as e:
+        return None
+
+
+def split_transcript(transcript, chunk_size=40):
+    sentences = []
+    for i in range(0, len(transcript), chunk_size):
+        to_add = [x["text"] for x in transcript[i : i + chunk_size]]
+        sentences.append(" ".join(to_add))
+    return sentences
+
+
+def query_punctuation(splits):
+    payload = {"inputs": splits}
+    API_URL = "https://api-inference.huggingface.co/models/oliverguhr/fullstop-punctuation-multilang-large"
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.json()
+
+
+def parse_output(output, comb):
+    total = []
+
+    # loop over the response from the huggingface api
+    for i, o in enumerate(output):
+        added = 0
+        tt = comb[i]
+        for elem in o:
+            # Loop over the output chunks and add the . and ?
+            if elem["entity_group"] not in ["0", ",", ""]:
+                split = elem["end"] + added
+                tt = tt[:split] + elem["entity_group"] + tt[split:]
+                added += 1
+        total.append(tt)
+    return " ".join(total)
+
+
+def punctuate(video_id):
+    transcript = retrieve_transcript(video_id)
+    splits = split_transcript(
+        transcript
+    )  # Get the transcript from the YoutubeTranscriptApi
+    resp = query_punctuation(splits)  # Get the response from the Inference API
+    punctuated_transcript = parse_output(resp, splits)
+    return punctuated_transcript