Update fn.py

fn.py

@@ -1,36 +1,48 @@
-from faster_whisper import WhisperModel
+import json
+import torch
+from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
+from datasets import load_dataset
+
+# config
+model_id = "kotoba-tech/kotoba-whisper-v1.0"
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 model = None
 model_size = None
+pipe = None
 
 def load_model(_model_size):
-    global model_size, model
+    global model_size, model, pipe
 
     if _model_size and model_size != _model_size:
         model_size = _model_size
 
-    try:
-        model = WhisperModel(model_size, device="cuda", compute_type="float16")
-    except:
-        model = WhisperModel(model_size, device="cpu", compute_type="int8")
+    # load model
+    model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True)
+    model.to(device)
+    processor = AutoProcessor.from_pretrained(model_id)
+    pipe = pipeline(
+        "automatic-speech-recognition",
+        model=model,
+        tokenizer=processor.tokenizer,
+        feature_extractor=processor.feature_extractor,
+        max_new_tokens=128,
+        torch_dtype=torch_dtype,
+        device=device,
+    )
 
 def speech_to_text(audio_file, _model_size = None):
-    global model_size, model
+    global model_size, model, pipe
 
     load_model(_model_size)
 
-    segments, info = model.transcribe(
-        audio_file,
-        language='ja',
-        beam_size=5,
-        vad_filter=True,
-        without_timestamps=False,
-    )
-
-    text_only = ''
-    text_with_timestamps = ''
-    for segment in segments:
-        text_only += f"{segment.text}\n"
-        text_with_timestamps += f"{segment.start:.2f}\t{segment.end:.2f}\t{segment.text}\n"
+    # run inference
+    result = pipe(audio_file)
 
-    return text_only, text_with_timestamps
+    try:
+        res = json.dumps(result)
+    except:
+        res = ''
+    
+    return result["text"], res