Views
No views yet
1def classify_query(text):
2 # Tokenize input
3 encoding = tokenizer(
4 text,
5 add_special_tokens=True,
6 max_length=128,
7 padding='max_length',
8 truncation=True,
9 return_attention_mask=True,
10 return_tensors='pt'
11 )
12
13 # Get prediction
14 model.eval()
15 with torch.no_grad():
16 outputs = model(
17 input_ids=encoding['input_ids'].to(device),
18 attention_mask=encoding['attention_mask'].to(device)
19 )
20
21 # Apply softmax to get probabilities
22 probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
23 class_0_prob = probs[0].item() # Not relevant probability
24 class_1_prob = probs[1].item() # Relevant probability
25
26 # Simple threshold-based classification
27 predicted_class = 1 if class_1_prob > 0.5 else 0
28
29 # Optional: Enhanced classification with keyword verification
30 laser_keywords = ["laser", "clean", "rust", "metal", "surface"]
31 contains_keywords = any(keyword in text.lower() for keyword in laser_keywords)
32
33 # Return classification result
34 if predicted_class == 1 or contains_keywords:
35 return "Relevant to laser cleaning"
36 else:
37 return "Not relevant to laser cleaning"