Update app.py

app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import pandas as pd
 import torch
 import torchaudio
-from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2Processor, safetensors
+from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2Processor
 from sklearn.preprocessing import LabelEncoder
 import numpy as np
 
@@ -10,7 +10,7 @@ import numpy as np
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load the fine-tuned model and processor
-model = Wav2Vec2ForSequenceClassification.from_pretrained("./fine_tuned_model", from_safetensors=True).to(device)
+model = Wav2Vec2ForSequenceClassification.from_pretrained("./fine_tuned_model", use_safetensors=True).to(device)
 processor = Wav2Vec2Processor.from_pretrained("./fine_tuned_model")
 
 # Load the label encoder
@@ -22,6 +22,7 @@ fixed_length = 10 * 16000  # 10 seconds * 16000 Hz
 
 # Function to get top 5 predictions with probabilities
 def get_top_5_predictions(logits, label_encoder):
+    logits = torch.tensor(logits)  # Convert numpy array to PyTorch tensor
     probabilities = torch.nn.functional.softmax(logits, dim=-1).cpu().numpy()
     top5_idx = np.argsort(probabilities, axis=-1)[:, -5:][:, ::-1]  # Top 5 indices
     top5_probs = np.take_along_axis(probabilities, top5_idx, axis=-1)
@@ -49,7 +50,7 @@ def predict(file_path):
     with torch.no_grad():
         logits = model(inputs.input_values).logits
     
-    return get_top_5_predictions(logits, label_encoder)
+    return get_top_5_predictions(logits.cpu().numpy(), label_encoder)
 
 # Streamlit interface
 st.title("Bird Sound Classification")