Spaces:
Running
Running
| 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 | |
| from pdf2image import convert_from_bytes | |
| import pytesseract | |
| from pytesseract import Output | |
| 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_scanned_pdf(self, file) -> str: | |
| """ | |
| Parameters: | |
| file (File): Scanned PDF file to be processed | |
| Returns: | |
| str: Text extracted from the scanned PDF file | |
| Extract text from scanned PDF file | |
| """ | |
| images = convert_from_bytes(file) | |
| all_text = "" | |
| for image in images: | |
| # Perform OCR on the image | |
| text = pytesseract.image_to_data(image, lang='eng', output_type=Output.DICT) | |
| # Extract text from the dictionary | |
| page_text = " ".join(text['text']) | |
| all_text += page_text + "\n" | |
| return all_text | |
| 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) | |