Added model prediction

app.py

@@ -1,22 +1,24 @@
 import json
 
 import streamlit as st
-from transformers import AutoTokenizer, RobertaForSequenceClassification
+from transformers import AutoTokenizer, RobertaForSequenceClassification, pipeline
 
 
 with open("config.json") as f:
     cfg = json.loads(f.read())
 
 
-# @st.cache(allow_output_mutation=True)
-# def load_model():
-#     tokenizer = AutoTokenizer.from_pretrained(cfg["model_name_or_path"])
-#     model = RobertaForSequenceClassification.from_pretrained(cfg["model_name_or_path"])
+@st.cache(allow_output_mutation=True)
+def load_model(input_text):
+    tokenizer = AutoTokenizer.from_pretrained(cfg["model_name_or_path"])
+    model = RobertaForSequenceClassification.from_pretrained(cfg["model_name_or_path"])
+
+    nlp = pipeline("text-classification", model=model, tokenizer=tokenizer)
+    result = nlp(input_text)
+    return result
 
-st.title("RoBERTa Marathi")
 
-# with st.spinner("Loading model..."):
-#     generator, tokenizer = load_model()
+st.title("RoBERTa Marathi")
 
 input_text = st.text_input("Text:")
 
@@ -25,4 +27,5 @@ predict_button = st.button("Predict")
 if predict_button:
     with st.spinner("Generating prediction..."):
         # Get prediction here
-        st.write("Predicted Label")
+        result = load_model(input_text)
+        st.write(result)