upload

README.md

@@ -0,0 +1,4 @@
+# Confidence Scoring
+
+Read https://x-lance.sjtu.edu.cn/papers/zhc00-chen-icassp17.pdf
+Run create_confidence_scores.py

create_confidence_scores.py

@@ -1,30 +1,57 @@
 #!/usr/bin/env python3
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
+from transformers import AutoModelForCTC, AutoProcessor
 from datasets import load_dataset
 import datasets
 import torch
+import sys
 
-model = Wav2Vec2ForCTC.from_pretrained("facebook/data2vec-audio-base-10m")
-processor = Wav2Vec2Processor.from_pretrained("facebook/data2vec-audio-base-10m")
+model_id = sys.argv[1]
 
-minds14 = load_dataset("PolyAI/minds14", "en-US", split="train")
-minds14 = minds14.cast_column("audio", datasets.Audio(sampling_rate=16_000))
+model = AutoModelForCTC.from_pretrained(model_id)
+processor = AutoProcessor.from_pretrained(model_id)
 
-input_values = processor(minds14[0]["audio"]["array"], return_tensors="pt", sampling_rate=minds14[0]["audio"]["sampling_rate"]).input_values
+num_samples = 4
+
+do_streaming = True
+
+if do_streaming:
+    dataset = load_dataset("common_voice", "en", split="test", streaming=True)
+    dataset = dataset.cast_column("audio", datasets.Audio(sampling_rate=16_000))
+
+    # iterate over dataset
+    dataset_iter = iter(dataset)
+    samples = [next(dataset_iter) for _ in range(num_samples)]
+
+    audio_samples = [s["audio"]["array"] for s in samples]
+    sampling_rate = set([s["audio"]["sampling_rate"] for s in samples]).pop()
+    text_samples = [s["sentence"] for s in samples]
+
+else:
+    dataset = load_dataset("patrickvonplaten/librispeech_asr_dummy", "clean", split="validation")
+    samples = dataset[:4]
+    audio_samples = [s["array"] for s in samples["audio"]]
+    sampling_rate = set([s["sampling_rate"] for s in samples["audio"]]).pop()
+    text_samples = samples["text"]
+
+inputs = processor(audio_samples, return_tensors="pt", sampling_rate=sampling_rate, padding=True)
 
 with torch.no_grad():
-    logits = model(input_values).logits
+    logits = model(inputs.input_values, attention_mask=inputs.attention_mask).logits
     scores = torch.nn.functional.softmax(logits, dim=-1)
     pred_ids = torch.argmax(logits, dim=-1)
     pred_scores = scores.gather(1, pred_ids.unsqueeze(-1))[:, :, 0]
 
 output = processor.batch_decode(pred_ids, output_word_offsets=True)
 
+
 # add confidence
-def confidence_score(word_dict):
-    probs = pred_scores[0, word_dict["start_offset"]: word_dict["end_offset"]]
-    return torch.mean(probs)
+def confidence_score(word_dict, index):
+    probs = pred_scores[index, word_dict["start_offset"]: word_dict["end_offset"]]
+    return round(torch.mean(probs).item(), 4)
 
-output["confidence_scores"] = {d["word"]: confidence_score(d) for d in output.word_offsets[0]}
 
-print(output["confidence_scores"])
+for i in range(num_samples):
+    print(20 * "=" + f"Output {i}" + 20 * "=")
+    print(text_samples[i])
+    print({d["word"]: confidence_score(d, i) for d in output.word_offsets[i]})
+    print("\n")