VladimirPozdeev/llm-security-scanner

1import re
2from transformers import AutoTokenizer,AutoModelForSequenceClassification
3def extract_response_and_analysis(full_response):
4
5    if '<|im_start|>assistant\n' in full_response:
6        full_response = full_response.split('<|im_start|>assistant\n')[-1]
7
8
9    full_response = full_response.replace('<|im_end|>', '').strip()
10
11    idx = full_response.rfind("Analysis:")
12
13    if idx == -1:
14        return full_response.strip(), None
15
16    response_part = full_response[:idx].strip()
17    analysis_part = full_response[idx:].strip()
18    return response_part, analysis_part
19
20def extract_analysis_fields(text):
21    result = {
22        'is_unsafe_prompt': None,
23        'attack_type': None,
24        'confidence': None,
25        'recommendation': None
26    }
27
28    patterns = {
29        'is_unsafe_prompt':      r'is_unsafe:\s*(\d+)',
30        'attack_type':    r'attack_type:\s*([^;]+?)(?:;|\n|$)',
31        'confidence':     r'confidence:\s*(high|medium|low)',
32        'recommendation': r'Recommendation:\s*(SAFE|REVIEW|BLOCK)',
33    }
34
35    for field, pattern in patterns.items():
36        match = re.search(pattern, str(text))
37        if match:
38            value = match.group(1).strip()
39            result[field] = int(value) if field == 'is_unsafe_prompt' else value
40
41    return result
42
43def load_pretrained_classification_model(path_to_model:str,device:str='cpu'):
44    tokenizer=AutoTokenizer.from_pretrained(path_to_model)
45    model=AutoModelForSequenceClassification.from_pretrained(path_to_model)
46    model.eval()
47    model.to(device)
48    return model, tokenizer