Spaces:
Running
Running
JuanjoSG5
commited on
Commit
•
688951f
1
Parent(s):
8f5ecb1
feat: added the main app
Browse files
app.py
ADDED
@@ -0,0 +1,61 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC, AutoTokenizer, BartForConditionalGeneration
|
3 |
+
import torch
|
4 |
+
import librosa
|
5 |
+
|
6 |
+
# Load BART tokenizer and model for summarization
|
7 |
+
tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
|
8 |
+
summarizer = BartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
|
9 |
+
|
10 |
+
# Load Wav2Vec2 processor and model for transcription
|
11 |
+
processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
|
12 |
+
model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
|
13 |
+
|
14 |
+
# Check if CUDA is available
|
15 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
16 |
+
model.to(device)
|
17 |
+
summarizer.to(device)
|
18 |
+
|
19 |
+
def transcribe_and_summarize(audioFile):
|
20 |
+
# Load audio as an array
|
21 |
+
audio, sampling_rate = librosa.load(audioFile, sr=16000) # Ensure it's 16kHz for Wav2Vec2
|
22 |
+
values = processor(audio, sampling_rate=sampling_rate, return_tensors="pt").input_values
|
23 |
+
|
24 |
+
# Move tensors to GPU if available
|
25 |
+
values = values.to(device)
|
26 |
+
|
27 |
+
# Transcription
|
28 |
+
with torch.no_grad():
|
29 |
+
logits = model(values).logits
|
30 |
+
predictedIDs = torch.argmax(logits, dim=-1)
|
31 |
+
transcription = processor.batch_decode(predictedIDs, skip_special_tokens=True)[0]
|
32 |
+
|
33 |
+
# Summarization
|
34 |
+
inputs = tokenizer(transcription, return_tensors="pt", truncation=True, max_length=1024)
|
35 |
+
inputs = inputs.to(device) # Move inputs to GPU
|
36 |
+
|
37 |
+
result = summarizer.generate(
|
38 |
+
inputs["input_ids"],
|
39 |
+
min_length=10,
|
40 |
+
max_length=256,
|
41 |
+
no_repeat_ngram_size=2,
|
42 |
+
encoder_no_repeat_ngram_size=2,
|
43 |
+
repetition_penalty=2.0,
|
44 |
+
num_beams=4,
|
45 |
+
early_stopping=True,
|
46 |
+
)
|
47 |
+
summary = tokenizer.decode(result[0], skip_special_tokens=True)
|
48 |
+
|
49 |
+
return transcription, summary
|
50 |
+
|
51 |
+
# Gradio interface
|
52 |
+
iface = gr.Interface(
|
53 |
+
fn=transcribe_and_summarize,
|
54 |
+
inputs=gr.Audio(type="filepath", label="Upload Audio"),
|
55 |
+
outputs=[gr.Textbox(label="Transcription"), gr.Textbox(label="Summary")],
|
56 |
+
title="Audio Transcription and Summarization",
|
57 |
+
description="Transcribe and summarize audio using Wav2Vec2 and BART.",
|
58 |
+
)
|
59 |
+
|
60 |
+
iface.launch()
|
61 |
+
|