Update README.md

README.md

@@ -13,25 +13,44 @@ import json
 import requests
 
 def convert_to_json(answer):
-    # Replace '<pad>' and '</s>' as per the provided logic
+    """
+    Convert a string representation of a dictionary to a JSON object.
+
+    This function takes a string representation of a dictionary, cleans it by removing
+    specific unwanted tokens and correcting boolean representations, and then converts
+    it into a JSON object.
+
+    Parameters:
+    answer (str): The input string representing a dictionary.
+
+    Returns:
+    dict: The JSON object converted from the input string.
+    """
     answer = answer.replace("<pad>", "").replace("</s>", "")
-    answer = answer.strip("'")  # Remove leading and trailing single quotes
-    
-    # Replace 'false' with 'False' and 'true' with 'True' for Python compatibility
+    answer = answer.strip("'")
     answer = answer.replace("false", "False").replace("true", "True")
-    
-    # Convert the string to a dictionary
     answer_dict = eval(answer)
-    
-    # Convert the dictionary to a JSON object
     answer_json = json.dumps(answer_dict)
-    
-    # Load the JSON object to ensure it's correctly formatted
     json_data = json.loads(answer_json)
-    
     return json_data
 
 def valid_url(url):
+    """
+    Validate the given URL against a list of supported platforms.
+
+    This function checks if the provided URL belongs to one of the supported
+    platforms for scanning. If the URL is valid, it returns True. Otherwise,
+    it returns a message indicating that the URL is not supported and lists the
+    available scanners.
+
+    Parameters:
+    url (str): The URL to be validated.
+
+    Returns:
+    bool or dict: Returns True if the URL is valid, otherwise returns a
+                  dictionary with a message indicating the URL is not supported
+                  and lists the available scanners.
+    """
     valid_list = [
         "github.com", "bitbucket.org", "sourceforge.net", "aws.amazon.com",
         "dev.azure.com", "gitea.com", "gogs.io", "phabricator.com",
@@ -41,8 +60,10 @@ def valid_url(url):
 
     if platform in valid_list:
         return True
-    else:
-        return {'message': 'provide an valid uri for scanning. currently we support pii_scanner, sast_scanner, sca_scanner, iac_scanner, container_scanner, malware_scanner, api_scanner'}
+
+    return {
+        'message': 'Provide a valid URL for scanning. Currently, we support PII_Scanner, SAST_Scanner, Sac_Scanner (Open_Source_Security), IaC_Scanner, Container_Scanner'
+    }
 ```
 
 ```python
@@ -50,8 +71,8 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 import time
 
-tokenizer = AutoTokenizer.from_pretrained("AquilaX-AI/NL-JSON-60M")
-model = AutoModelForSeq2SeqLM.from_pretrained("AquilaX-AI/NL-JSON-60M")
+tokenizer = AutoTokenizer.from_pretrained("AquilaX-AI/NL-JSON-Start-Scan")
+model = AutoModelForSeq2SeqLM.from_pretrained("AquilaX-AI/NL-JSON-Start-Scan")
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -67,15 +88,20 @@ inputs = inputs.to(device)
 outputs = model.generate(**inputs, max_length=256)
 answer = tokenizer.decode(outputs[0])
 json_data = convert_to_json(answer)
+
+to_return = json_data.copy()
 try:
     valid = valid_url(json_data["repo"])
-    if valid == True:
-        print(json_data)
+    if valid != True:
+        to_return = valid
     else:
-        print(valid)
+      url = re.findall(r'https?://\S+', query)
+      to_return['repo'] = url
+
 except:
-    print(json_data)
+    pass
 
 end = time.time()
+print(to_return)
 print(f"Time taken: {end - start}")
 ```