File size: 4,018 Bytes
0aaead9 | 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 | """
Example inference script for AIBA BERT-BiLSTM model
This script demonstrates how to use the model for prediction.
"""
import torch
from transformers import AutoTokenizer
from nn_model import load_model_and_tokenizer
def predict(text, model_path="YOUR_USERNAME/aiba-bert-bilstm"):
"""
Make predictions on input text
Args:
text: Input text to analyze
model_path: Hugging Face model repo ID
Returns:
dict with extracted entities, intent, and language
"""
# Load model
model, tokenizer, config = load_model_and_tokenizer(model_path)
# Get label mappings from config
id2tag = {int(k): v for k, v in config['id2tag'].items()}
id2intent = {int(k): v for k, v in config['id2intent'].items()}
id2lang = {int(k): v for k, v in config['id2lang'].items()}
# Tokenize
inputs = tokenizer(
text,
return_tensors="pt",
truncation=True,
max_length=config['max_length'],
return_offsets_mapping=True
)
offset_mapping = inputs.pop('offset_mapping')[0]
# Predict
with torch.no_grad():
outputs = model(**inputs)
# Process NER predictions
ner_predictions = torch.argmax(outputs['ner_logits'], dim=2)[0]
tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
# Extract entities
entities = []
current_entity = None
current_tokens = []
for token, pred_id, offset in zip(tokens, ner_predictions, offset_mapping):
if token in ['[CLS]', '[SEP]', '[PAD]']:
continue
tag = id2tag[pred_id.item()]
if tag.startswith('B-'):
if current_entity:
entity_text = text[current_tokens[0][0]:current_tokens[-1][1]]
entities.append({'type': current_entity, 'value': entity_text})
current_entity = tag[2:]
current_tokens = [offset.tolist()]
elif tag.startswith('I-') and current_entity:
current_tokens.append(offset.tolist())
else:
if current_entity:
entity_text = text[current_tokens[0][0]:current_tokens[-1][1]]
entities.append({'type': current_entity, 'value': entity_text})
current_entity = None
current_tokens = []
if current_entity:
entity_text = text[current_tokens[0][0]:current_tokens[-1][1]]
entities.append({'type': current_entity, 'value': entity_text})
# Process intent prediction
intent_logits = outputs['intent_logits'][0]
intent_probs = torch.nn.functional.softmax(intent_logits, dim=0)
intent_id = torch.argmax(intent_probs).item()
intent = id2intent[intent_id]
intent_confidence = intent_probs[intent_id].item()
# Process language prediction
lang_logits = outputs['lang_logits'][0]
lang_probs = torch.nn.functional.softmax(lang_logits, dim=0)
lang_id = torch.argmax(lang_probs).item()
language = id2lang[lang_id]
lang_confidence = lang_probs[lang_id].item()
return {
'text': text,
'entities': entities,
'intent': intent,
'intent_confidence': intent_confidence,
'language': language,
'language_confidence': lang_confidence
}
# Example usage
if __name__ == "__main__":
# Test examples
examples = [
"Qabul qiluvchi Omad Biznes MCHJ STIR 123456789 summa 500000 UZS",
"Получатель ООО Прогресс ИНН 987654321 сумма 1000000 руб",
"Transfer 5000 USD to Starlight Ventures LLC TIN 555666777"
]
for text in examples:
print(f"\nInput: {text}")
result = predict(text)
print(f"Intent: {result['intent']} ({result['intent_confidence']:.2%})")
print(f"Language: {result['language']} ({result['language_confidence']:.2%})")
print("Entities:")
for entity in result['entities']:
print(f" - {entity['type']}: {entity['value']}")
|