Update app.py

app.py

@@ -1,17 +1,23 @@
 import streamlit as st
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from transformers import RobertaTokenizer, RobertaModel
+from your_model_file import MDFEND  # Ensure you import your model correctly
 
+# Load model and tokenizer
 @st.cache(allow_output_mutation=True)
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("./prediction_sinhala.ipynb")
-    model = AutoModelForSequenceClassification.from_pretrained("./prediction_sinhala.ipynb")
+    tokenizer = RobertaTokenizer.from_pretrained("prediction_sinhala.ipynb")
+    model = MDFEND.from_pretrained("prediction_sinhala.ipynb")
     return model, tokenizer
 
 model, tokenizer = load_model()
 
+# User input
 text_input = st.text_area("Enter text here:")
+
+# Prediction
 if st.button("Predict"):
     inputs = tokenizer(text_input, return_tensors="pt")
-    outputs = model(**inputs)
+    with torch.no_grad():  # Ensure no gradients are computed
+        outputs = model(**inputs)
     prediction = outputs.logits.argmax(-1).item()
     st.write(f"Prediction: {prediction}")