Views
No views yet
1from transformers import AutoModelForTokenClassification, AutoTokenizer
2import torch
3
4model = AutoModelForTokenClassification.from_pretrained("MichaelHuang/muril_base_cased_urdu_ner")
5tokenizer = AutoTokenizer.from_pretrained("google/muril-base-cased")
6
7# Define the labels dictionary
8labels_dict = {
9 0: "B-FESTIVAL",
10 1: "B-GAME",
11 2: "B-LANGUAGE",
12 3: "B-LITERATURE",
13 4: "B-LOCATION",
14 5: "B-MISC",
15 6: "B-NUMEX",
16 7: "B-ORGANIZATION",
17 8: "B-PERSON",
18 9: "B-RELIGION",
19 10: "B-TIMEX",
20 11: "I-FESTIVAL",
21 12: "I-GAME",
22 13: "I-LANGUAGE",
23 14: "I-LITERATURE",
24 15: "I-LOCATION",
25 16: "I-MISC",
26 17: "I-NUMEX",
27 18: "I-ORGANIZATION",
28 19: "I-PERSON",
29 20: "I-RELIGION",
30 21: "I-TIMEX",
31 22: "O"
32}
33
34def ner_predict(sentence, model, tokenizer, labels_dict):
35 # Tokenize the input sentence
36 inputs = tokenizer(sentence, return_tensors="pt", padding=True, truncation=True, max_length=128)
37
38 # Perform inference
39 with torch.no_grad():
40 outputs = model(**inputs)
41
42 # Get the predicted labels
43 predicted_labels = torch.argmax(outputs.logits, dim=2)
44
45 # Convert tokens and labels to lists
46 tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
47 labels = predicted_labels.squeeze().tolist()
48
49 # Map numeric labels to string labels
50 predicted_labels = [labels_dict[label] for label in labels]
51
52 # Combine tokens and labels
53 result = list(zip(tokens, predicted_labels))
54
55 return result
56
57test_sentence = "امیتابھ اور ریکھا کی فلم 'گنگا کی سوگندھ' 10 فروری سنہ 1978 کو ریلیز ہوئی تھی۔ اس کے بعد راکھی، رندھیر کپور اور نیتو سنگھ کے ساتھ 'قسمے وعدے' 21 اپریل 1978 کو ریلیز ہوئی۔"
58predictions = ner_predict(test_sentence, model, tokenizer, labels_dict)
59
60for token, label in predictions:
61 print(f"{token}: {label}")