File size: 7,812 Bytes
99d1428 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 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.") |