Views
No views yet
1import re
2import urllib.parse
3
4from transformers import AutoTokenizer, AutoModelForSequenceClassification
5import nltk.tokenize
6import torch
7
8preprocess_tokenizer_regex = r'[^\W_0-9]+|[^\w\s]+|_+|\s+|[0-9]+' # Similar to wordpunct_tokenize
9preprocess_tokenizer = nltk.tokenize.RegexpTokenizer(preprocess_tokenizer_regex).tokenize
10
11def preprocess_url(url):
12 protocol_idx = url.find("://")
13 protocol_idx = (protocol_idx + 3) if protocol_idx != -1 else 0
14 url = url.rstrip('/')[protocol_idx:]
15 url = urllib.parse.unquote(url, errors="backslashreplace")
16
17 # Remove blanks
18 url = re.sub(r'\s+', ' ', url)
19 url = re.sub(r'^\s+|\s+$', '', url)
20
21 # Tokenize
22 url = ' '.join(preprocess_tokenizer(url))
23
24 return url
25
26tokenizer = AutoTokenizer.from_pretrained("Transducens/xlm-roberta-base-parallel-urls-classifier")
27model = AutoModelForSequenceClassification.from_pretrained("Transducens/xlm-roberta-base-parallel-urls-classifier")
28
29# prepare input
30url1 = preprocess_url("https://web.ua.es/en/culture.html")
31url2 = preprocess_url("https://web.ua.es/es/cultura.html")
32urls = f"{url1}{tokenizer.sep_token}{url2}"
33encoded_input = tokenizer(urls, add_special_tokens=True, truncation=True, padding="longest",
34 return_attention_mask=True, return_tensors="pt", max_length=256)
35
36# forward pass
37output = model(encoded_input["input_ids"], encoded_input["attention_mask"])
38
39# obtain probability
40probability = torch.sigmoid(output["logits"]).cpu().squeeze().item()
41
42print(probability)