forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
import json
|
|
import re
|
|
import warnings
|
|
from typing import List, Tuple
|
|
|
|
import requests
|
|
|
|
from langchain.docstore.document import Document
|
|
from langchain.document_loaders.base import BaseLoader
|
|
|
|
|
|
class BiliBiliLoader(BaseLoader):
|
|
"""Loader that loads bilibili transcripts."""
|
|
|
|
def __init__(self, video_urls: List[str]):
|
|
"""Initialize with bilibili url."""
|
|
self.video_urls = video_urls
|
|
|
|
def load(self) -> List[Document]:
|
|
"""Load from bilibili url."""
|
|
results = []
|
|
for url in self.video_urls:
|
|
transcript, video_info = self._get_bilibili_subs_and_info(url)
|
|
doc = Document(page_content=transcript, metadata=video_info)
|
|
results.append(doc)
|
|
|
|
return results
|
|
|
|
def _get_bilibili_subs_and_info(self, url: str) -> Tuple[str, dict]:
|
|
try:
|
|
from bilibili_api import sync, video
|
|
except ImportError:
|
|
raise ValueError(
|
|
"requests package not found, please install it with "
|
|
"`pip install bilibili-api-python`"
|
|
)
|
|
|
|
bvid = re.search(r"BV\w+", url)
|
|
if bvid is not None:
|
|
v = video.Video(bvid=bvid.group())
|
|
else:
|
|
aid = re.search(r"av[0-9]+", url)
|
|
if aid is not None:
|
|
try:
|
|
v = video.Video(aid=int(aid.group()[2:]))
|
|
except AttributeError:
|
|
raise ValueError(f"{url} is not bilibili url.")
|
|
else:
|
|
raise ValueError(f"{url} is not bilibili url.")
|
|
|
|
video_info = sync(v.get_info())
|
|
video_info.update({"url": url})
|
|
|
|
# Get subtitle url
|
|
subtitle = video_info.pop("subtitle")
|
|
sub_list = subtitle["list"]
|
|
if sub_list:
|
|
sub_url = sub_list[0]["subtitle_url"]
|
|
result = requests.get(sub_url)
|
|
raw_sub_titles = json.loads(result.content)["body"]
|
|
raw_transcript = " ".join([c["content"] for c in raw_sub_titles])
|
|
|
|
raw_transcript_with_meta_info = (
|
|
f"Video Title: {video_info['title']},"
|
|
f"description: {video_info['desc']}\n\n"
|
|
f"Transcript: {raw_transcript}"
|
|
)
|
|
return raw_transcript_with_meta_info, video_info
|
|
else:
|
|
raw_transcript = ""
|
|
warnings.warn(
|
|
f"""
|
|
No subtitles found for video: {url}.
|
|
Return Empty transcript.
|
|
"""
|
|
)
|
|
return raw_transcript, video_info
|