Spaces:
Runtime error
Runtime error
import pickle | |
import pandas as pd | |
import gradio as gr | |
import numpy as np | |
from sentence_transformers import SentenceTransformer, util | |
from transformers import pipeline, Wav2Vec2ProcessorWithLM | |
from librosa import load, resample | |
# Constants | |
model_name = 'sentence-transformers/msmarco-distilbert-base-v4' | |
max_sequence_length = 512 | |
# Load corpus | |
import subprocess | |
subprocess.run(["gdown", "1QVpyk_xyqNYrHT3NdUfBxbDV_eyCDa2Q"]) | |
with open("embeddings.pkl", "rb") as fp: | |
pickled_data = pickle.load(fp) | |
sentences = pickled_data['sentences'] | |
corpus_embeddings = pickled_data['embeddings'] | |
print(f'Number of documents: {len(sentences)}') | |
# Load pre-embedded corpus | |
print(f'Number of embeddings: {corpus_embeddings.shape[0]}') | |
# Load embedding model | |
model = SentenceTransformer(model_name) | |
model.max_seq_length = max_sequence_length | |
# Load speech to text model | |
asr_model = "patrickvonplaten/wav2vec2-base-960h-4-gram" | |
processor = Wav2Vec2ProcessorWithLM.from_pretrained(asr_model) | |
asr = pipeline( | |
"automatic-speech-recognition", | |
model=asr_model, | |
tokenizer=processor.tokenizer, | |
feature_extractor=processor.feature_extractor, | |
decoder=processor.decoder, | |
) | |
def find_sentences(query, hits): | |
query_embedding = model.encode(query) | |
hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=hits) | |
hits = hits[0] | |
output_texts = [] | |
output_scores = [] | |
for hit in hits: | |
# Find source document based on sentence index | |
output_texts.append(sentences[hit['corpus_id']]) | |
output_scores.append(hit['score']) | |
return pd.DataFrame(data={"Text": output_texts, "Score": output_scores}) | |
def process(input_selection, query, filepath, hits): | |
if input_selection=='speech': | |
speech, sampling_rate = load(filepath) | |
if sampling_rate != 16000: | |
speech = resample(speech, sampling_rate, 16000) | |
text = asr(speech)['text'] | |
else: | |
text = query | |
return text, find_sentences(text, hits) | |
# Gradio inputs | |
buttons = gr.inputs.Radio(['text','speech'], type='value', default='speech', label='Input selection') | |
text_query = gr.inputs.Textbox(lines=1, label='Text input', default='breast cancer biomarkers') | |
mic = gr.inputs.Audio(source='microphone', type='filepath', label='Speech input', optional=True) | |
slider = gr.inputs.Slider(minimum=1, maximum=10, step=1, default=3, label='Number of hits') | |
# Gradio outputs | |
speech_query = gr.outputs.Textbox(type='auto', label='Query string') | |
results = gr.outputs.Dataframe( | |
headers=['Text', 'Score'], | |
col_width=200, | |
label='Query results') | |
iface = gr.Interface( | |
theme='huggingface', | |
description='This Space lets you query a text corpus containing 50,000 random clinical trial descriptions', | |
fn=process, | |
layout='horizontal', | |
inputs=[buttons,text_query,mic,slider], | |
outputs=[speech_query, results], | |
examples=[ | |
['text', "breast cancer biomarkers", 'dummy.wav', 3], | |
], | |
allow_flagging=False | |
) | |
iface.launch() | |