video-search / handler.py

use device_number for setting GPU

010ad9e over 1 year ago

No virus

10.3 kB

	"""
	https://huggingface.co/tomiwa1a/video-search
	"""
	from typing import Dict

	from sentence_transformers import SentenceTransformer
	from tqdm import tqdm
	import whisper
	from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
	import torch
	import pytube
	import time


	class EndpointHandler():
	# load the model
	WHISPER_MODEL_NAME = "tiny.en"
	SENTENCE_TRANSFORMER_MODEL_NAME = "multi-qa-mpnet-base-dot-v1"
	QUESTION_ANSWER_MODEL_NAME = "vblagoje/bart_lfqa"
	SUMMARIZER_MODEL_NAME = "philschmid/bart-large-cnn-samsum"
	device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
	device_number = 0 if torch.cuda.is_available() else -1

	def __init__(self, path=""):

	device = "cuda" if torch.cuda.is_available() else "cpu"
	device_number = 0 if torch.cuda.is_available() else -1
	print(f'whisper and question_answer_model will use: {device}')
	print(f'whisper and question_answer_model will use device_number: {device_number}')

	t0 = time.time()
	self.whisper_model = whisper.load_model(self.WHISPER_MODEL_NAME).to(device_number)
	t1 = time.time()

	total = t1 - t0
	print(f'Finished loading whisper_model in {total} seconds')

	t0 = time.time()
	self.sentence_transformer_model = SentenceTransformer(self.SENTENCE_TRANSFORMER_MODEL_NAME)
	t1 = time.time()

	total = t1 - t0
	print(f'Finished loading sentence_transformer_model in {total} seconds')

	t0 = time.time()
	self.summarizer = pipeline("summarization", model=self.SUMMARIZER_MODEL_NAME, device=device_number)
	t1 = time.time()

	total = t1 - t0
	print(f'Finished loading summarizer in {total} seconds')

	self.question_answer_tokenizer = AutoTokenizer.from_pretrained(self.QUESTION_ANSWER_MODEL_NAME)
	t0 = time.time()
	self.question_answer_model = AutoModelForSeq2SeqLM.from_pretrained \
	(self.QUESTION_ANSWER_MODEL_NAME).to(device_number)
	t1 = time.time()
	total = t1 - t0
	print(f'Finished loading question_answer_model in {total} seconds')

	def __call__(self, data: Dict[str, str]) -> Dict:
	"""
	Args:
	data (:obj:):
	includes the URL to video for transcription
	Return:
	A :obj:`dict`:. transcribed dict
	"""
	# process input
	print('data', data)

	if "inputs" not in data:
	raise Exception(f"data is missing 'inputs' key which EndpointHandler expects. Received: {data}"
	f" See: https://huggingface.co/docs/inference-endpoints/guides/custom_handler#2-create-endpointhandler-cp")
	video_url = data.pop("video_url", None)
	query = data.pop("query", None)
	long_form_answer = data.pop("long_form_answer", None)
	summarize = data.pop("summarize", False)
	encoded_segments = {}
	if video_url:
	video_with_transcript = self.transcribe_video(video_url)
	video_with_transcript['transcript']['transcription_source'] = f"whisper_{self.WHISPER_MODEL_NAME}"
	encode_transcript = data.pop("encode_transcript", True)
	if encode_transcript:
	encoded_segments = self.combine_transcripts(video_with_transcript)
	encoded_segments = {
	"encoded_segments": self.encode_sentences(encoded_segments)
	}
	return {
	**video_with_transcript,
	**encoded_segments
	}
	elif summarize:
	summary = self.summarize_video(data["segments"])
	return {"summary": summary}
	elif query:
	if long_form_answer:
	context = data.pop("context", None)
	answer = self.generate_answer(query, context)
	response = {
	"answer": answer
	}

	return response
	else:
	query = [{"text": query, "id": ""}] if isinstance(query, str) else query
	encoded_segments = self.encode_sentences(query)

	response = {
	"encoded_segments": encoded_segments
	}

	return response

	else:
	return {
	"error": "'video_url' or 'query' must be provided"
	}

	def transcribe_video(self, video_url):
	decode_options = {
	# Set language to None to support multilingual,
	# but it will take longer to process while it detects the language.
	# Realized this by running in verbose mode and seeing how much time
	# was spent on the decoding language step
	"language": "en",
	"verbose": True
	}
	yt = pytube.YouTube(video_url)
	video_info = {
	'id': yt.video_id,
	'thumbnail': yt.thumbnail_url,
	'title': yt.title,
	'views': yt.views,
	'length': yt.length,
	# Althhough, this might seem redundant since we already have id
	# but it allows the link to the video be accessed in 1-click in the API response
	'url': f"https://www.youtube.com/watch?v={yt.video_id}"
	}
	stream = yt.streams.filter(only_audio=True)[0]
	path_to_audio = f"{yt.video_id}.mp3"
	stream.download(filename=path_to_audio)
	t0 = time.time()
	transcript = self.whisper_model.transcribe(path_to_audio, **decode_options)
	t1 = time.time()
	for segment in transcript['segments']:
	# Remove the tokens array, it makes the response too verbose
	segment.pop('tokens', None)

	total = t1 - t0
	print(f'Finished transcription in {total} seconds')

	# postprocess the prediction
	return {"transcript": transcript, 'video': video_info}

	def encode_sentences(self, transcripts, batch_size=64):
	"""
	Encoding all of our segments at once or storing them locally would require too much compute or memory.
	So we do it in batches of 64
	:param transcripts:
	:param batch_size:
	:return:
	"""
	# loop through in batches of 64
	all_batches = []
	for i in tqdm(range(0, len(transcripts), batch_size)):
	# find end position of batch (for when we hit end of data)
	i_end = min(len(transcripts), i + batch_size)
	# extract the metadata like text, start/end positions, etc
	batch_meta = [{
	**row
	} for row in transcripts[i:i_end]]
	# extract only text to be encoded by embedding model
	batch_text = [
	row['text'] for row in batch_meta
	]
	# create the embedding vectors
	batch_vectors = self.sentence_transformer_model.encode(batch_text).tolist()

	batch_details = [
	{
	**batch_meta[x],
	'vectors': batch_vectors[x]
	} for x in range(0, len(batch_meta))
	]
	all_batches.extend(batch_details)

	return all_batches

	def summarize_video(self, segments):
	for index, segment in enumerate(segments):
	segment['summary'] = self.summarizer(segment['text'])
	segment['summary'] = segment['summary'][0]['summary_text']
	print('index', index)
	print('length', segment['length'])
	print('text', segment['text'])
	print('summary', segment['summary'])

	return segments

	def generate_answer(self, query, documents):

	# concatenate question and support documents into BART input
	conditioned_doc = "<P> " + " <P> ".join([d for d in documents])
	query_and_docs = "question: {} context: {}".format(query, conditioned_doc)

	model_input = self.question_answer_tokenizer(query_and_docs, truncation=False, padding=True,
	return_tensors="pt")

	generated_answers_encoded = self.question_answer_model.generate(
	input_ids=model_input["input_ids"].to(self.device),
	attention_mask=model_input["attention_mask"].to(self.device),
	min_length=64,
	max_length=256,
	do_sample=False,
	early_stopping=True,
	num_beams=8,
	temperature=1.0,
	top_k=None,
	top_p=None,
	eos_token_id=self.question_answer_tokenizer.eos_token_id,
	no_repeat_ngram_size=3,
	num_return_sequences=1)
	answer = self.question_answer_tokenizer.batch_decode(generated_answers_encoded, skip_special_tokens=True,
	clean_up_tokenization_spaces=True)
	return answer

	@staticmethod
	def combine_transcripts(video, window=6, stride=3):
	"""

	:param video:
	:param window: number of sentences to combine
	:param stride: number of sentences to 'stride' over, used to create overlap
	:return:
	"""
	new_transcript_segments = []

	video_info = video['video']
	transcript_segments = video['transcript']['segments']
	for i in tqdm(range(0, len(transcript_segments), stride)):
	i_end = min(len(transcript_segments), i + window)
	text = ' '.join(transcript['text']
	for transcript in
	transcript_segments[i:i_end])
	# TODO: Should int (float to seconds) conversion happen at the API level?
	start = int(transcript_segments[i]['start'])
	end = int(transcript_segments[i]['end'])
	new_transcript_segments.append({
	**video_info,
	**{
	'start': start,
	'end': end,
	'title': video_info['title'],
	'text': text,
	'id': f"{video_info['id']}-t{start}",
	'url': f"https://youtu.be/{video_info['id']}?t={start}",
	'video_id': video_info['id'],
	}
	})
	return new_transcript_segments