Update handler.py

handler.py

@@ -2,18 +2,19 @@ from typing import  List, Any
 import torch
 from peft import PeftModel, PeftConfig
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from youtube_transcript_api import YouTubeTranscriptApi
 
 def preprocessing(data):          
     texts = list()
     i = 0
-    if len(data) <= i+4000:
+    if len(data) <= i+3000:
         texts = data
     else:
         while len(data[i:]) != 0:
-            if len(data[i:]) > 4000:
-                string = str(data[i:i+4000])
+            if len(data[i:]) > 3000:
+                string = str(data[i:i+3000])
                 texts.append(string)
-                i = i + 3800
+                i = i + 2800
             else:
                 string = str(data[i:])
                 texts.append(string)
@@ -33,10 +34,16 @@ class EndpointHandler:
 
     def __call__(self, data: Any) -> List[str]:
         
-        inputs = data.pop("inputs", data)
+        video_id = data.pop("inputs", data)
+        dict = YouTubeTranscriptApi.get_transcript(video_id)
+        
+        transcript = ""
+
+        for i in range(len(dict)):
+            transcript += dict[i]['text']
 
         # process input
-        texts = preprocessing(inputs)
+        texts = preprocessing(transcript)
         inputs = self.tokenizer(texts, return_tensors="pt", padding=True, ) # truncation=True
         
         with torch.no_grad():