model update

main.py

@@ -1,35 +1,12 @@
 from flask import Flask, request, jsonify
-from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
 import numpy as np
 import torch
 
 app = Flask(__name__)
-model = AutoModelForSpeechSeq2Seq.from_pretrained("GitNazarov/whisper-small-pt-3-uz")
-processor = AutoProcessor.from_pretrained("GitNazarov/whisper-small-pt-3-uz")
-
-
-USE_ONNX = False # change this to True if you want to test onnx model
-silero_vad_path = 'snakers4/silero-vad'
-vad_model, vad_utils = torch.hub.load(silero_vad_path,
-                              model='silero_vad',
-                              force_reload=True,
-                              onnx=USE_ONNX)
-
-(get_speech_timestamps,
-save_audio,
-read_audio,
-VADIterator,
-collect_chunks) = vad_utils
-STT_SAMPLE_RATE = 16000
-
-
-def int2float(sound):
-    abs_max = np.abs(sound).max()
-    sound = sound.astype('float32')
-    if abs_max > 0:
-        sound *= 1/32768
-    sound = sound.squeeze()  # depends on the use case
-    return sound
+processor = Wav2Vec2Processor.from_pretrained("oyqiz/uzbek_stt")
+model = Wav2Vec2ForCTC.from_pretrained("oyqiz/uzbek_stt")
+SAMPLE_RATE = 16000
 
 @app.route('/', methods=['GET'])
 def index():
@@ -38,21 +15,16 @@ def index():
 @app.route('/transcribe', methods=['POST'])
 def transcribe():
     data_frames = request.data
-    audio_data = np.frombuffer(data_frames, dtype=np.int16)
-    audio_float = int2float(audio_data)
-    final_data = torch.from_numpy(audio_float)
-    sp_timestamps = get_speech_timestamps(final_data, vad_model, sampling_rate=STT_SAMPLE_RATE)
-    try:
-        final_audio_data = collect_chunks(sp_timestamps, final_data)
-        inputs = processor(final_audio_data, return_tensors="pt", sampling_rate=16000, max_new_tokens=100)
-        input_features = inputs.input_features
-        generated_ids = model.generate(inputs=input_features)
+    audio_np = np.frombuffer(data_frames, dtype=np.int16)
+    audio_np = audio_np / np.iinfo(np.int16).max
+    inputs = processor(audio_np, sampling_rate=SAMPLE_RATE, return_tensors="pt")
+
+    with torch.no_grad():
+        logits = model(inputs.input_values, attention_mask=inputs.attention_mask).logits
 
-        transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)
-        transcription = ''.join(transcription)
-    except Exception as e:
-        transcription = str(e)
-    return str(transcription), {'Content-Type': 'application/json'}
+    predicted_ids = torch.argmax(logits, dim=-1)
+    transcription = processor.decode(predicted_ids[0])
+    return transcription
 
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860)