Spaces:
Sleeping
Sleeping
File size: 6,396 Bytes
a828a8b 9bf726b a828a8b 9bf726b a828a8b 9bf726b a828a8b 9bf726b a828a8b 9bf726b a828a8b 9bf726b a828a8b 9bf726b a828a8b 9bf726b a828a8b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 |
import os
import json
from typing import List
import assemblyai as aai
from youtube_transcript_api import YouTubeTranscriptApi
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.documents.base import Document
from langchain_community.vectorstores import DeepLake
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.prompts.few_shot import FewShotPromptTemplate
from langchain.prompts.prompt import PromptTemplate
from typing import Dict
import uuid
class Processing:
def __init__(self, dataset_path: str, embedding_model_name: str, chunk_size=500, chunk_overlap=5):
"""
Parameters:
dataset_path (str): Path to the dataset in the Vector-DB
embedding_model_name (str): Name of the HuggingFace model to be used for embeddings
chunk_size (int): Size of each chunk to be processed
chunk_overlap (int): Overlap between each chunk
Initialize embedding model, text splitter, transcriber and Vector-DB
"""
self.dataset_path = dataset_path
self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
self.transcriber = aai.Transcriber()
self.embedding_model = HuggingFaceEmbeddings(
model_name=embedding_model_name,
encode_kwargs={'normalize_embeddings': False}
)
self.db = DeepLake(dataset_path=f"hub://{self.dataset_path}",
embedding=self.embedding_model,
exec_option="compute_engine"
)
def _add_metadata(self, documents: List[Document], url: str, id: str, source: str, file_type: str,
course_tag="") -> (List[Document], Dict[str, str]):
"""
Parameters:
documents (List[Document]): List of documents to add metadata to
id (str): ID of the documents
source (str): Source of the documents
file_type (str): Type of the documents
course_tag (str): Tag to identify the course the documents belongs to
Returns:
documents (List[Document], Dict[str, str): List of documents with metadata added along with the metadata
Add metadata to the documents
"""
metadata = {
"id": id,
"source": source,
"url": url,
"file_type": file_type,
"course_tag": course_tag
}
for doc in documents:
doc.metadata = metadata
return documents, metadata
def load_pdf(self, text) -> (List[Document], Dict[str, str]):
"""
Returns:
documents (List[Document], Dict[str, str): List of documents with metadata added along with the metadata
Load PDF file, split into chunks and add metadata
"""
pdf_chunk = self.text_splitter.create_documents([text])
print("Created document chunks")
return self._add_metadata(pdf_chunk, url="NaN", id=str(uuid.uuid4()), source="document", file_type="pdf")
def load_transcript(self, url) -> (List[Document], Dict[str, str]):
"""
Returns:
documents (List[Document], Dict[str, str): List of documents with metadata added along with the metadata
Load transcript, split into chunks and add metadata
"""
transcript = self.transcriber.transcribe(url)
print("Transcribed")
transcript_chunk = self.text_splitter.create_documents([transcript.text])
print("Created transcript chunks")
return self._add_metadata(transcript_chunk, url="NaN", id=str(uuid.uuid4()), source="custom_video",
file_type="transcript")
def load_yt_transcript(self, url) -> (List[Document], Dict[str, str]):
"""
Returns:
documents (List[Document], Dict[str, str): List of documents with metadata added along with the metadata
Load YouTube transcript, split into chunks and add metadata
"""
if url.startswith("https://www.youtube.com/watch?v="):
video_id = url.replace("https://www.youtube.com/watch?v=", "")
else:
video_id = url.replace("https://youtu.be/", "")
transcript = YouTubeTranscriptApi.get_transcript(video_id)
print("Downloaded transcript")
transcript = [line['text'] for line in transcript]
transcript_text = ' '.join(transcript)
yt_transcript_chunk = self.text_splitter.create_documents([transcript_text])
print("Created YouTube transcript chunks")
return self._add_metadata(yt_transcript_chunk, url=url, id=video_id, source="youtube", file_type="transcript")
def upload_to_db(self, documents: List[Document]):
"""
Parameters:
documents (List[Document]): List of documents to upload to Vector-DB
Upload documents to Vector-DB
"""
print("Embedding and uploading to Vector-DB...")
self.db.add_documents(documents)
print("Uploaded to Vector-DB")
class PromptCreate:
def __init__(self, example_path, save_path):
self.examples = []
self.example_prompt = None
self.few_shot_prompt = None
self.example_path = example_path
self.file_name = "example_{i}.json"
self.save_path = save_path
def load_examples(self):
for i in range(1, len(os.listdir(self.example_path)) + 1):
filename = os.path.join(self.example_path, self.file_name.format(i=i))
try:
with open(filename, "r") as json_file:
self.examples.append(json.load(json_file))
except FileNotFoundError:
print(f"File {filename} not found.")
except json.JSONDecodeError:
print(f"Error decoding JSON from file {filename}.")
def create_prompt_template(self, input_variables, template_string):
self.example_prompt = PromptTemplate(input_variables=input_variables, template=template_string)
def create_few_shot_prompt(self, prefix, suffix):
self.few_shot_prompt = FewShotPromptTemplate(
examples=self.examples, example_prompt=self.example_prompt, prefix=prefix, suffix=suffix
)
def save_prompt(self):
self.few_shot_prompt.save(self.save_path)
|