Update app.py

app.py

@@ -12,8 +12,14 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # Load the processor
 processor = Wav2Vec2Processor.from_pretrained("./fine_tuned_model")
 
-# Load the entire model
-model = torch.load("./fine_tuned_model/model.pt", map_location=device)
+# Initialize the model
+model = Wav2Vec2ForSequenceClassification.from_pretrained(
+    "facebook/wav2vec2-base-960h",
+    num_labels=len(pd.read_csv("dataset/train_wav.csv")["Common Name"].unique())
+)
+
+# Load the model's state dictionary
+model.load_state_dict(torch.load("./fine_tuned_model/model_state_dict.pt", map_location=device))
 model.to(device)
 model.eval()
 
@@ -75,4 +81,4 @@ if uploaded_file is not None:
             top5_predictions = predict(file_path)
             st.success("Top 5 Predicted Bird Species with Probabilities:")
             for label, prob in top5_predictions:
-                st.write(f"{label}: {prob:.4f}")
+                st.write(f"{label}: {prob:.4f}")