chgrdj commited on
Commit
64aa5df
·
verified ·
1 Parent(s): fa5fc06

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +40 -0
app.py ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
3
+ from peft import PeftModel
4
+ import torch
5
+
6
+ @st.cache_resource
7
+ def load_model():
8
+ model_id = "google/flan-t5-large"
9
+ adapter_path = "./Flan-T5-Typosquat-detect" # Adjust to your saved adapter path
10
+
11
+ # Load the tokenizer and model
12
+ tokenizer = AutoTokenizer.from_pretrained(model_id)
13
+ model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
14
+ model = PeftModel.from_pretrained(model, adapter_path)
15
+ model.eval()
16
+
17
+ return model, tokenizer
18
+
19
+ model, tokenizer = load_model()
20
+
21
+ st.title("FLAN-T5 Typosquatting Detection")
22
+ st.write("Enter a potential typosquatted domain and a target domain to check if one is a variant of the other.")
23
+
24
+ potential_typosquat = st.text_input("Potential Typosquatted Domain", value="lonlonsoft.com")
25
+ target_domain = st.text_input("Target Domain", value="stiltsoft.net")
26
+
27
+ if st.button("Check Typosquatting"):
28
+ if potential_typosquat and target_domain:
29
+
30
+ prompt = f"Is the first domain a typosquat of the second: {potential_typosquat} {target_domain}"
31
+
32
+ input_ids = tokenizer(prompt, return_tensors="pt").input_ids
33
+ outputs = model.generate(input_ids, max_new_tokens=20)
34
+
35
+ prediction = tokenizer.decode(outputs[0], skip_special_tokens=True)
36
+
37
+ st.write("**Prediction:**")
38
+ st.write(prediction)
39
+ else:
40
+ st.warning("Please enter both domains to perform the check.")