|
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
|
|
|
|
|
|
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'
|
|
|
|
|
|
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)
|
|
|
|
|
|
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.") |
