Update main.py

main.py

@@ -105,19 +105,24 @@ async def upload_audio(file: UploadFile = File(...)):
     print("length of the audio array:" , len(audio))
 
     print("*" * 100)
-        
-    # Tokenize audio
+
+    # Tokenization
     print("Tokenizing audio...")
-    input_values = tokenizer(audio, return_tensors="pt").input_values
-    
+    input_values = await asyncio.to_thread(tokenizer, audio, return_tensors="pt")
+    input_values = input_values.input_values
+    print("Tokenization complete. Shape of input_values:", input_values.shape)
+
     # Perform inference
     print("Performing inference with Wav2Vec2 model...")
-    logits = model(input_values).logits
-    
+    output = await asyncio.to_thread(model, input_values)
+    logits = output.logits
+    print("Inference complete. Shape of logits:", logits.shape)
+
     # Get predictions
     print("Getting predictions...")
     prediction = torch.argmax(logits, dim=-1)
-    
+    print("Prediction shape:", prediction.shape)
+
     # Decode predictions
     print("Decoding predictions...")
     transcription = tokenizer.batch_decode(prediction)[0]
@@ -125,8 +130,8 @@ async def upload_audio(file: UploadFile = File(...)):
     # Convert transcription to lowercase
     transcription = transcription.lower()
     
-    # Print transcription and word counts
     print("Decoded transcription:", transcription)
+    
     incorrect, correct = count_spelled_words(transcription, english_words)
     print("Spelling check - Incorrect words:", incorrect, ", Correct words:", correct)