Views
No views yet
pip install -U transformers[torch], then copy the snippet from the section.pip install transformers[torch]1import json
2import requests
3
4def convert_to_json(answer):
5 """
6 Convert a string representation of a dictionary to a JSON object.
7
8 This function takes a string representation of a dictionary, cleans it by removing
9 specific unwanted tokens and correcting boolean representations, and then converts
10 it into a JSON object.
11
12 Parameters:
13 answer (str): The input string representing a dictionary.
14
15 Returns:
16 dict: The JSON object converted from the input string.
17 """
18 answer = answer.replace("<pad>", "").replace("</s>", "")
19 answer = answer.strip("'")
20 answer = answer.replace("false", "False").replace("true", "True")
21 answer_dict = eval(answer)
22 answer_json = json.dumps(answer_dict)
23 json_data = json.loads(answer_json)
24 return json_data
25
26def valid_url(url):
27 """
28 Validate the given URL against a list of supported platforms.
29
30 This function checks if the provided URL belongs to one of the supported
31 platforms for scanning. If the URL is valid, it returns True. Otherwise,
32 it returns a message indicating that the URL is not supported and lists the
33 available scanners.
34
35 Parameters:
36 url (str): The URL to be validated.
37
38 Returns:
39 bool or dict: Returns True if the URL is valid, otherwise returns a
40 dictionary with a message indicating the URL is not supported
41 and lists the available scanners.
42 """
43 valid_list = [
44 "github.com", "bitbucket.org", "sourceforge.net", "aws.amazon.com",
45 "dev.azure.com", "gitea.com", "gogs.io", "phabricator.com",
46 "gitkraken.com", "beanstalkapp.com", "gitlab.com"
47 ]
48 platform = url.split("//")[1].split("/")[0]
49
50 if platform in valid_list:
51 return True
52
53 return {
54 'message': 'Provide a valid URL for scanning. Currently, we support PII_Scanner, SAST_Scanner, Sac_Scanner (Open_Source_Security), IaC_Scanner, Container_Scanner'
55 }1from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
2import torch
3import time
4
5tokenizer = AutoTokenizer.from_pretrained("AquilaX-AI/NL-JSON-Start-Scan")
6model = AutoModelForSeq2SeqLM.from_pretrained("AquilaX-AI/NL-JSON-Start-Scan")
7
8device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
9
10# Change YOUR_QUERY eg: can this https://github.com/mr-vicky-01/educational-assitant on every week using pii and sast scan
11query = "Translate the following text to JSON: " + "YOUR_QUERY".lower()
12query = query.replace(",", "")
13
14start = time.time()
15
16inputs = tokenizer(query, return_tensors="pt")
17model.to(device)
18inputs = inputs.to(device)
19outputs = model.generate(**inputs, max_length=256)
20answer = tokenizer.decode(outputs[0])
21try:
22 json_data = convert_to_json(answer)
23except:
24 json_data = {'message': 'We encountered an issue with your query. Please use the Personalized Scan option for accurate results.'}
25
26to_return = json_data.copy()
27try:
28 valid = valid_url(json_data["repo"])
29 if valid != True:
30 to_return = valid
31 else:
32 url = re.findall(r'https?://\S+', query)
33 to_return['repo'] = url
34
35except:
36 pass
37
38end = time.time()
39print(to_return)
40print(f"Time taken: {end - start}")