Update app.py

app.py

@@ -15,26 +15,38 @@ def load_model():
     model.eval()
     
     return model, tokenizer
-
+device='cpu'
 model, tokenizer = load_model()
 
 st.title("FLAN-T5 Typosquatting Detection")
 st.write("Enter a potential typosquatted domain and a target domain to check if one is a variant of the other.")
 
+# Non-editable prompt part
+prompt_prefix = "Is the first domain a typosquat of the second:"
+
+# Display the non-editable prompt with input fields for the rest
+st.markdown("### Prompt")
+st.text_area("Prompt", prompt_prefix, height=50, disabled=True)
+
+# User inputs for dynamic part of the prompt
 potential_typosquat = st.text_input("Potential Typosquatted Domain", value="lonlonsoft.com")
-target_domain = st.text_input("Target Domain", value="stiltsoft.net")
+target_domain = st.text_input("Legitimate Domain", value="stiltsoft.net")
 
+# Generate prompt by combining fixed and dynamic parts
+full_prompt = f"{prompt_prefix} {potential_typosquat} {target_domain}"
+
+# Perform inference when button is clicked
 if st.button("Check Typosquatting"):
     if potential_typosquat and target_domain:
-        
-        prompt = f"Is the first domain a typosquat of the second: {potential_typosquat} {target_domain}"
-
-        input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+        # Encode and generate response
+        input_ids = tokenizer(full_prompt, return_tensors="pt").input_ids.to(device)
         outputs = model.generate(input_ids, max_new_tokens=20)
         
+        # Decode the response
         prediction = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
+        # Display the result
         st.write("**Prediction:**")
         st.write(prediction)
     else:
-        st.warning("Please enter both domains to perform the check.")
+        st.warning("Please enter both domains to perform the check.")