Update app.py

app.py

@@ -1,3 +1,51 @@
-import gradio as gr
+# Load model directly
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import numpy as np
+import torch
 
-gr.load("models/kmack/malicious-url-detection").launch()
+# Check if CUDA is available
+if torch.cuda.is_available():
+    # Choose a specific GPU or use the default
+    device = torch.device("cuda:0")
+else:
+# Or CPU
+    device = torch.device("cpu")
+
+tokenizer = AutoTokenizer.from_pretrained("kmack/malicious-url-detection")
+model = AutoModelForSequenceClassification.from_pretrained("kmack/malicious-url-detection")
+
+# set Model to cude
+model = model.to(device)
+
+# predict function
+def get_predit(input_text: str) -> dict:
+    label2id = model.config.label2id
+    inputs = tokenizer(input_text, return_tensors='pt', truncation=True)
+    inputs = inputs.to(device)
+    outputs = model(**inputs)
+    logits = outputs.logits
+    sigmoid = torch.nn.Sigmoid()
+    probs = sigmoid(logits.squeeze().cpu())
+    probs = probs.detach().numpy()
+    for i, k in enumerate(label2id.keys()):
+        label2id[k] = probs[i]
+    label2id = {k: float(v) for k, v in sorted(label2id.items(), key=lambda item: item[1].item(), reverse=True)}
+    return label2id
+
+# Define example URLs
+example_url_1 = 'https://medium.com'
+example_url_2 = 'http://google.com-redirect@valimail.com'
+example_url_3 = 'https://a101-nisan-kampanyalari.com'
+
+# Create the Gradio interface
+demo = gr.Interface(
+    fn=get_predit,
+    inputs=gr.components.Textbox(label='Input', placeholder='Enter URL here...'),
+    outputs=gr.components.Label(label='Predictions', num_top_classes=5),
+    title='kmack/malicious-url-detection',
+    description='Detects whether a given URL is benign or potentially malicious.',
+    examples=[[example_url_1], [example_url_2], [example_url_3]],
+    allow_flagging='never'
+)
+
+demo.launch()