Views
No views yet

1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2import torch
3
4
5model_name = "epfl-dhlab/CatastroBERT"
6tokenizer = AutoTokenizer.from_pretrained(model_name)
7model = AutoModelForSequenceClassification(model_name)
8
9def predict(text):
10 # Prepare the text data
11 inputs = tokenizer.encode_plus(
12 text,
13 None,
14 add_special_tokens=True,
15 return_token_type_ids=True,
16 padding=True,
17 max_length=512,
18 truncation=True,
19 return_tensors='pt'
20 )
21
22 ids = inputs['input_ids'].to('cuda' if torch.cuda.is_available() else 'cpu')
23 mask = inputs['attention_mask'].to('cuda' if torch.cuda.is_available() else 'cpu')
24
25 # Get predictions
26 with torch.no_grad():
27 outputs = model(ids, mask)
28 logits = outputs.logits
29
30 # Apply sigmoid function to get probabilities
31 probs = torch.sigmoid(logits).cpu().numpy()
32
33 # Return the probability of the class (1)
34 return probs[0][0]
35
36#example usage
37text = "Un violent ouragan du sud-ouest est passé cette nuit sur Lausanne."
38print(f"Prediction: {predict(text)}")