Update app.py

app.py

@@ -1,40 +1,48 @@
-# app.py
+# app.py (Updated Version)
 
-# Import the necessary libraries
 import gradio as gr
 from transformers import pipeline
+import pickle
 
 # =============================================================================
-# 1. LOAD YOUR MODEL
+# 1. LOAD YOUR MODEL AND THE SAVED LABEL ENCODER
 # =============================================================================
-# Use a pipeline for easy text classification. This will automatically
-# load your fine-tuned model and tokenizer from the repository.
-#
-# IMPORTANT: Replace "your-username/your-model-repo-name" with your actual model path.
-# For example: "Tarive/esm2_t12_35M_UR50D-finetuned-pfam-1k"
-#
-model_path = "Tarive/esm2_t12_35M_UR50D-finetuned-pfam-1k"
+# Define the path to your model repository
+model_path = "Tarive/esm2_t12_35M_UR50D-finetuned-pfam-1k" # Make sure this is correct
+
+# Load the classification pipeline
 classifier = pipeline("text-classification", model=model_path)
 
+# Load the label encoder from the file you uploaded
+with open("label_encoder.pkl", "rb") as f:
+    label_encoder = pickle.load(f)
+
 
 # =============================================================================
-# 2. DEFINE THE PREDICTION FUNCTION
+# 2. DEFINE THE PREDICTION FUNCTION WITH LABEL DECODING
 # =============================================================================
-# This function takes a sequence as input and returns a formatted result.
+# This function now decodes the labels before displaying them.
 def predict_family(sequence):
-    # The pipeline returns a list of dictionaries.
-    # Example: [{'label': 'PF00042', 'score': 0.979}]
-    predictions = classifier(sequence, top_k=5) # Get the top 5 predictions
+    # Get the top 5 predictions from the model
+    predictions = classifier(sequence, top_k=5)
     
-    # Format the results into a more readable dictionary for display.
-    results = {p['label']: p['score'] for p in predictions}
+    # The model outputs labels like "LABEL_455". We need to extract the number.
+    results = {}
+    for p in predictions:
+        # Extract the number from the label string (e.g., "LABEL_455" -> 455)
+        label_index = int(p['label'].split('_')[1])
+        
+        # Use the label_encoder to find the original family name
+        original_label = label_encoder.inverse_transform([label_index])[0]
+        
+        # Store the real name and score
+        results[original_label] = p['score']
     
     return results
 
 # =============================================================================
-# 3. CREATE THE GRADIO INTERFACE
+# 3. CREATE THE GRADIO INTERFACE (No changes here)
 # =============================================================================
-# This creates the actual web page interface.
 iface = gr.Interface(
     fn=predict_family,
     inputs=gr.Textbox(
@@ -52,7 +60,7 @@ iface = gr.Interface(
         ["MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTKTYFPHFDLSHGSAQVKGHGKKVADALTNAVAHVDDMPNALSALSDLHAHKLRVDPVNFKLLSHCLLVTLAAHLPAEFTPAVHASLDKFLASVSTVLTSKYR"],
         ["MTEYKLVVVGAGDVGKSALTIQLIQNHFVDEYDPTIEDSYRKQVEVDCQQCMILDILDTAGQEEYSAMRDQYMRTGEGFLCVFAINNTKSFEDIHQYREQIKRVKDSDDVPMVLVGNKCDLAARTVESRQAQDLARSYGIPYIETSAKTRQGVEDAFYTLVREIRQHKLRKLNPPDESGGCMS"]
     ],
-    allow_flagging="never" # Disables the "Flag" button for a cleaner interface
+    allow_flagging="never"
 )
 
 # Launch the interface!