wehub-resource-sync
98 lines
3.1 KiB
Python
98 lines
3.1 KiB
Python
from datetime import datetime
|
|
import pixeltable as pxt
|
|
from pixeltable.functions import openai
|
|
from pixeltable.functions.huggingface import sentence_transformer
|
|
from pixeltable.functions.video import extract_audio
|
|
from pixeltable.iterators import AudioSplitter, FrameIterator
|
|
from pixeltable.iterators.string import StringSplitter
|
|
from pixeltable.functions.openai import vision
|
|
|
|
EMBED_MODEL = sentence_transformer.using(model_id='intfloat/e5-large-v2')
|
|
|
|
# Set to True to delete existing index
|
|
directory = 'video_index'
|
|
table_name = f'{directory}.video'
|
|
|
|
# Create video table
|
|
pxt.create_dir(directory, if_exists='replace_force')
|
|
video_index = pxt.create_table(table_name, {'video': pxt.Video, 'uploaded_at': pxt.Timestamp})
|
|
video_index.add_computed_column(audio_extract=extract_audio(video_index.video, format='mp3'))
|
|
|
|
# Create view for frames
|
|
frames_view = pxt.create_view(
|
|
f'{directory}.video_frames',
|
|
video_index,
|
|
iterator=FrameIterator.create(
|
|
video=video_index.video,
|
|
fps=1
|
|
)
|
|
)
|
|
|
|
frames_view.add_computed_column(
|
|
image_description=vision(
|
|
prompt="Provide quick caption for the image.",
|
|
image=frames_view.frame,
|
|
model="gpt-4o-mini"
|
|
)
|
|
)
|
|
|
|
frames_view.add_embedding_index('image_description', string_embed=EMBED_MODEL)
|
|
|
|
# Create view for audio chunks
|
|
chunks_view = pxt.create_view(
|
|
f'{directory}.video_chunks',
|
|
video_index,
|
|
iterator=AudioSplitter.create(
|
|
audio=video_index.audio_extract,
|
|
chunk_duration_sec=30.0,
|
|
overlap_sec=2.0,
|
|
min_chunk_duration_sec=5.0
|
|
)
|
|
)
|
|
|
|
# Audio-to-text for chunks
|
|
chunks_view.add_computed_column(
|
|
transcription=openai.transcriptions(audio=chunks_view.audio_chunk, model='whisper-1')
|
|
)
|
|
|
|
# Create view that chunks text into sentences
|
|
transcription_chunks = pxt.create_view(
|
|
f'{directory}.video_sentence_chunks',
|
|
chunks_view,
|
|
iterator=StringSplitter.create(text=chunks_view.transcription.text, separators='sentence'),
|
|
)
|
|
|
|
# Create embedding index
|
|
transcription_chunks.add_embedding_index('text', string_embed=EMBED_MODEL)
|
|
|
|
# Insert Videos
|
|
videos = [
|
|
'https://github.com/pixeltable/pixeltable/raw/release/docs/resources/audio-transcription-demo/'
|
|
f'Lex-Fridman-Podcast-430-Excerpt-{n}.mp4'
|
|
for n in range(3)
|
|
]
|
|
|
|
video_index.insert({'video': video, 'uploaded_at': datetime.now()} for video in videos[:2])
|
|
|
|
# Get similarity scores
|
|
audio_sim = transcription_chunks.text.similarity('What is happiness?')
|
|
image_sim = frames_view.image_description.similarity('Black Suit')
|
|
|
|
# Fetch 5 most similar audio chunks
|
|
audio_results = (
|
|
transcription_chunks.order_by(audio_sim, transcription_chunks.uploaded_at, asc=False)
|
|
.limit(5)
|
|
.select(transcription_chunks.text, transcription_chunks.uploaded_at, similarity=audio_sim)
|
|
.collect()
|
|
)
|
|
|
|
# Fetch 5 most similar frames
|
|
frame_results = (
|
|
frames_view.order_by(image_sim, frames_view.uploaded_at, asc=False)
|
|
.limit(5)
|
|
.select(frames_view.image_description, frames_view.uploaded_at, similarity=image_sim)
|
|
.collect()
|
|
)
|
|
|
|
print(audio_results)
|
|
print(frame_results) |