wehub-resource-sync
68 lines
2.5 KiB
Python
68 lines
2.5 KiB
Python
import pixeltable as pxt
|
|
from pixeltable.functions import whisper
|
|
from pixeltable.functions.huggingface import sentence_transformer
|
|
from pixeltable.iterators.string import StringSplitter
|
|
from pixeltable.iterators import AudioSplitter
|
|
|
|
DIRECTORY = 'audio_index'
|
|
TABLE_NAME = f'{DIRECTORY}.audio'
|
|
CHUNKS_VIEW_NAME = f'{DIRECTORY}.audio_chunks'
|
|
SENTENCES_VIEW_NAME = f'{DIRECTORY}.audio_sentence_chunks'
|
|
DELETE_INDEX = True
|
|
|
|
# python -m spacy download en_core_web_sm (run this separately if needed)
|
|
|
|
if DELETE_INDEX and TABLE_NAME in pxt.list_tables():
|
|
pxt.drop_table(TABLE_NAME, force=True)
|
|
|
|
if TABLE_NAME not in pxt.list_tables():
|
|
# Create audio table
|
|
pxt.create_dir(DIRECTORY, if_exists='ignore')
|
|
audio_index = pxt.create_table(TABLE_NAME, {'audio_file': pxt.Audio})
|
|
|
|
# Create view for audio chunks
|
|
chunks_view = pxt.create_view(
|
|
CHUNKS_VIEW_NAME,
|
|
audio_index,
|
|
iterator=AudioSplitter.create(
|
|
audio=audio_index.audio_file,
|
|
chunk_duration_sec=30.0, # Split into 30-second chunks
|
|
overlap_sec=2.0, # 2-second overlap between chunks
|
|
min_chunk_duration_sec=5.0 # Drop last chunk if < 5 seconds
|
|
)
|
|
)
|
|
|
|
# Create audio-to-text column on chunks
|
|
chunks_view.add_computed_column(
|
|
transcription=whisper.transcribe(audio=chunks_view.audio_chunk, model='base.en')
|
|
)
|
|
|
|
# Create view that chunks text into sentences
|
|
sentences_view = pxt.create_view(
|
|
SENTENCES_VIEW_NAME,
|
|
chunks_view,
|
|
iterator=StringSplitter.create(text=chunks_view.transcription.text, separators='sentence'),
|
|
)
|
|
|
|
# Define the embedding model
|
|
embed_model = sentence_transformer.using(model_id='intfloat/e5-large-v2')
|
|
|
|
# Create embedding index
|
|
sentences_view.add_embedding_index(column='text', string_embed=embed_model)
|
|
else:
|
|
audio_index = pxt.get_table(TABLE_NAME)
|
|
chunks_view = pxt.get_view(CHUNKS_VIEW_NAME)
|
|
sentences_view = pxt.get_view(SENTENCES_VIEW_NAME)
|
|
|
|
# Add data to the table
|
|
audio_index.insert([{'audio_file': 's3://pixeltable-public/audio/10-minute tour of Pixeltable.mp3'}])
|
|
|
|
# Semantic search
|
|
query_text = 'What is Pixeltable?'
|
|
|
|
# Calculate similarity scores between query and sentences
|
|
sim = sentences_view.text.similarity(query_text)
|
|
|
|
# Get top 5 most similar sentences with their scores
|
|
results = sentences_view.order_by(sim, asc=False).select(sentences_view.text, sim=sim).limit(5).collect()
|
|
print(results['text']) |