Update app.py

app.py

@@ -12,6 +12,7 @@ def load_model():
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
     model = PeftModel.from_pretrained(model, adapter_path)
+    model = model.merge_and_unload()
     model.eval()
     
     return model, tokenizer
@@ -21,21 +22,13 @@ 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=68, 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("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:
         # Encode and generate response
@@ -46,7 +39,7 @@ if st.button("Check Typosquatting"):
         prediction = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
         # Display the result
-        st.write("**Prediction:**")
+        st.write("**Prediction: **")
         st.write(prediction)
     else:
         st.warning("Please enter both domains to perform the check.")