Upload app.py

app.py

@@ -0,0 +1,166 @@
+import streamlit as st
+import torch
+from transformers import AlbertTokenizer, AlbertForSequenceClassification, AlbertModel
+import numpy as np
+import pandas as pd
+import os
+from torch.nn.functional import softmax
+import torch.nn as nn
+
+# Paths
+LEVEL_DIRS = {
+    1: 'level1',
+    2: 'level2',
+    3: 'level3',
+    4: 'level4',
+    5: 'level5',
+    6: 'level6',
+    7: 'level7'
+}
+MAPPING_FILE = 'mapping.csv'
+MODEL_NAME = 'albert/albert-base-v2'  # Define the base model name
+
+# Load mapping
+mapping_df = pd.read_csv(MAPPING_FILE)
+
+def get_label_text(level, predicted_id):
+    level_map = {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6}
+    level_num = level_map.get(level)
+    if level_num is not None:
+        row = mapping_df[(mapping_df['level'] == level_num) & (mapping_df['id'] == predicted_id)]
+        return row['text'].iloc[0] if not row.empty else "Description not found"
+    return "Invalid Level"
+
+def predict_level(level, text, parent_prediction_id=None, checkpoint_path=None):
+    level_dir = LEVEL_DIRS[level]
+    tokenizer = AlbertTokenizer.from_pretrained(checkpoint_path)
+    label_map = np.load(os.path.join(level_dir, 'label_map.npy'), allow_pickle=True).item()
+    num_labels = len(label_map)
+
+    if level == 1:
+        model = AlbertForSequenceClassification.from_pretrained(checkpoint_path)
+    else:
+        parent_level_dir = LEVEL_DIRS[level - 1]
+        parent_label_map = np.load(os.path.join(parent_level_dir, 'label_map.npy'), allow_pickle=True).item()
+        num_parent_labels = len(parent_label_map)
+
+        class TaxonomyClassifier(nn.Module):
+            def __init__(self, base_model_name, num_parent_labels, num_labels):
+                super().__init__()
+                self.albert = AlbertModel.from_pretrained(base_model_name)
+                self.dropout = nn.Dropout(0.1)
+                self.classifier = nn.Linear(self.albert.config.hidden_size + num_parent_labels, num_labels)
+
+            def forward(self, input_ids, attention_mask, parent_ids):
+                outputs = self.albert(input_ids, attention_mask=attention_mask)
+                pooled_output = outputs.pooler_output
+                pooled_output = self.dropout(pooled_output)
+                combined_features = torch.cat((pooled_output, parent_ids), dim=1)
+                logits = self.classifier(combined_features)
+                return logits
+
+        model = TaxonomyClassifier(MODEL_NAME, num_parent_labels, num_labels)
+        model.load_state_dict(torch.load(os.path.join(checkpoint_path, 'model.safetensors'), map_location=torch.device('cpu')))
+
+    model.eval()
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+
+    if level > 1:
+        parent_label_map_current = np.load(os.path.join(LEVEL_DIRS[level - 1], 'label_map.npy'), allow_pickle=True).item()
+        num_parent_labels_current = len(parent_label_map_current)
+        parent_one_hot = torch.zeros(num_parent_labels_current)
+        if parent_prediction_id != 0:
+            parent_index = parent_label_map_current.get(parent_prediction_id)
+            if parent_index is not None:
+                parent_one_hot[parent_index] = 1.0
+        with torch.no_grad():
+            outputs = model(inputs.input_ids, attention_mask=inputs.attention_mask, parent_ids=parent_one_hot.unsqueeze(0))
+    else:
+        with torch.no_grad():
+            outputs = model(**inputs)
+
+    probabilities = softmax(outputs.logits if level == 1 else outputs, dim=-1)[0]
+    top3_prob, top3_indices = torch.topk(probabilities, 3)
+    index_to_label = {v: k for k, v in label_map.items()}
+    results = []
+    for prob, index in zip(top3_prob, top3_indices):
+        predicted_label_id = index_to_label[index.item()]
+        results.append((predicted_label_id, prob.item()))
+    return results
+
+st.title("Taxonomy Model Inference")
+
+input_text = st.text_area("Enter text to classify", "Experience the magic of music with the Clavinova CLP-800 series. This versatile range of digital pianos is designed to delight everyone, from budding musicians to seasoned pianists. Each model combines state-of-the-art technology with the realistic touch and tone of world-renowned grand pianos, enhanced by GrandTouch keyboard action and Virtual Resonance Modeling. With seamless Bluetooth® connectivity, built-in lessons, and elegant design, the CLP-800 series offers the perfect blend of tradition and innovation. Elevate your musical journey with the warmth and sophistication of the Yamaha Clavinova, our finest series of digital pianos.")
+
+softmax_threshold = st.slider("Softmax Threshold", min_value=0.0, max_value=1.0, value=0.5, step=0.05)
+
+# Checkpoint Selection
+available_levels = []
+level_checkpoints = {}
+for level in LEVEL_DIRS:
+    level_dir = LEVEL_DIRS[level]
+    if os.path.exists(level_dir):
+        options = [d for d in os.listdir(level_dir) if os.path.isdir(os.path.join(level_dir, d))]
+        options = [d for d in options if 'step' in d or d == 'model']
+        options.sort(key=lambda x: (('step' not in x), int(x.split('step')[-1]) if 'step' in x else -1))
+        level_checkpoints[level] = [os.path.join(level_dir, opt) for opt in options]
+        if level_checkpoints[level]:
+            available_levels.append(level)
+    else:
+        level_checkpoints[level] = []
+
+selected_checkpoints = {}
+for level in available_levels:
+    selected_checkpoints[level] = st.selectbox(f"Select Level {level} Checkpoint", options=level_checkpoints[level])
+
+if st.button("Run Inference"):
+    if input_text:
+        all_level_results = {}
+        current_prediction_id = None
+        last_level = 0
+
+        for level in sorted(available_levels):
+            if selected_checkpoints[level]:
+                checkpoint_path = selected_checkpoints[level]
+                if level == 1:
+                    level_results = predict_level(level, input_text, checkpoint_path=checkpoint_path)
+                else:
+                    if current_prediction_id == 0:
+                        st.info(f"Taxonomy terminated at Level {last_level} with ID 0.")
+                        break
+                    level_results = predict_level(level, input_text, parent_prediction_id=current_prediction_id, checkpoint_path=checkpoint_path)
+
+                if level_results[0][1] < softmax_threshold:
+                    st.info(f"Inference stopped at Level {level} due to softmax probability ({level_results[0][1]:.3f}) being below the threshold.")
+                    break
+
+                all_level_results[level] = level_results
+                current_prediction_id = level_results[0][0]
+                last_level = level
+            else:
+                st.warning(f"Skipping Level {level} as no checkpoint is selected.")
+                break
+
+        data = []
+        for level in sorted(all_level_results.keys()):
+            results = all_level_results[level]
+            data.append({
+                'level': level,
+                'text': get_label_text(level - 1, results[0][0]),
+                'softmax': f"{results[0][1]:.3f}",
+                'runner_up_1_id': results[1][0],
+                'runner_up_1_text': get_label_text(level - 1, results[1][0]),
+                'runner_up_1_softmax': f"{results[1][1]:.3f}",
+                'runner_up_2_id': results[2][0],
+                'runner_up_2_text': get_label_text(level - 1, results[2][0]),
+                'runner_up_2_softmax': f"{results[2][1]:.3f}",
+            })
+
+        if data:
+            df = pd.DataFrame(data)
+            st.dataframe(df)
+        else:
+            st.info("No predictions made or inference stopped.")
+
+    else:
+        st.warning("Please enter text for classification.")