Update SampleQA.py

SampleQA.py

@@ -1,75 +1,66 @@
-from os import path
-import streamlit as st
-import tensorflow as tf
-import random
-from transformers import ElectraTokenizerFast, TFElectraForQuestionAnswering
-from datasets import Dataset, DatasetDict, load_dataset
-
-model_hf = "nguyennghia0902/bestfailed_electra-small-discriminator_5e-05_16"
-tokenizer = ElectraTokenizerFast.from_pretrained(model_hf)
-reload_model = TFElectraForQuestionAnswering.from_pretrained(model_hf)
-
-@st.cache_resource
-def predict(question, context):
-    inputs = tokenizer(question, context, return_offsets_mapping=True,return_tensors="tf",max_length=512, truncation=True)
-    offset_mapping = inputs.pop("offset_mapping")
-    outputs = reload_model(**inputs)
-    answer_start_index = int(tf.math.argmax(outputs.start_logits, axis=-1)[0])
-    answer_end_index = int(tf.math.argmax(outputs.end_logits, axis=-1)[0])
-    start_char = offset_mapping[0][answer_start_index][0]
-    end_char = offset_mapping[0][answer_end_index][1]
-    predicted_answer_text = context[start_char:end_char]
-    
-    return predicted_answer_text
-
-def main():
-    st.set_page_config(page_title="Sample in Dataset", page_icon="📝")
-
-    # giving a title to our page
-    col1, col2 = st.columns([2, 1])
-    col1.title("Sample in Dataset")
-
-    new_data = load_dataset("nguyennghia0902/project02_textming_dataset", data_files={'train': 'raw_newformat_data/traindata-00000-of-00001.arrow', 'test': 'raw_newformat_data/testdata-00000-of-00001.arrow'})
-    
-    sampleQ = ""
-    sampleC = ""
-    sampleA = ""
-    if st.button("Sample"):
-        sample = random.choice(new_data['test'])
-        sampleQ = sample['question']
-        sampleC = sample['context']
-        sampleA = sample['answers']["text"][0]
-
-    question = st.text_area(
-        "Sample QUESTION: ",
-        sampleQ,
-        height=15,
-    )
-    text = st.text_area(
-        "Sample CONTEXT:",
-        sampleC,
-        height=100,
-    )
-    answer = st.text_area(
-        "True ANSWER:",
-        sampleA,
-        height=20,
-    )
-
-    # Create a prediction button
-    if st.button("Predict"):
-        prediction = ""
-        stripped_text = text.strip()
-        if not stripped_text:
-            st.error("Please enter a context.")
-            return
-        stripped_question = question.strip()
-        if not stripped_question:
-            st.error("Please enter a question.")
-            return
-
-        prediction = predict(stripped_question, stripped_text)
-        st.success(prediction)
-
-if __name__ == "__main__":
+from os import path
+import streamlit as st
+import tensorflow as tf
+import random
+from transformers import ElectraTokenizerFast, TFElectraForQuestionAnswering
+from datasets import Dataset, DatasetDict, load_dataset
+
+model_hf = "nguyennghia0902/bestfailed_electra-small-discriminator_5e-05_16"
+tokenizer = ElectraTokenizerFast.from_pretrained(model_hf)
+reload_model = TFElectraForQuestionAnswering.from_pretrained(model_hf)
+
+@st.cache_resource
+def predict(question, context):
+    inputs = tokenizer(question, context, return_offsets_mapping=True,return_tensors="tf",max_length=512, truncation=True)
+    offset_mapping = inputs.pop("offset_mapping")
+    outputs = reload_model(**inputs)
+    answer_start_index = int(tf.math.argmax(outputs.start_logits, axis=-1)[0])
+    answer_end_index = int(tf.math.argmax(outputs.end_logits, axis=-1)[0])
+    start_char = offset_mapping[0][answer_start_index][0]
+    end_char = offset_mapping[0][answer_end_index][1]
+    predicted_answer_text = context[start_char:end_char]
+    
+    return predicted_answer_text
+
+def main():
+    st.set_page_config(page_title="Sample in Dataset", page_icon="📝")
+
+    # giving a title to our page
+    col1, col2 = st.columns([2, 1])
+    col1.title("Sample in Dataset")
+
+    new_data = load_dataset("nguyennghia0902/project02_textming_dataset", data_files={'train': 'raw_newformat_data/traindata-00000-of-00001.arrow', 'test': 'raw_newformat_data/testdata-00000-of-00001.arrow'})
+    
+    sampleQ = ""
+    sampleC = ""
+    sampleA = ""
+    if st.button("Sample"):
+        sample = random.choice(new_data['test'])
+        sampleQ = sample['question']
+        sampleC = sample['context']
+        sampleA = sample['answers']["text"][0]
+
+    question = st.text_area(
+        "Sample QUESTION: ",
+        sampleQ,
+        height=15,
+    )
+    text = st.text_area(
+        "Sample CONTEXT:",
+        sampleC,
+        height=100,
+    )
+    answer = st.text_area(
+        "True ANSWER:",
+        sampleA,
+        height=20,
+    )
+
+    # Create a prediction button
+    if st.button("Predict"):
+        prediction = ""
+        prediction = predict(sampleQ, sampleC)
+        st.success(prediction)
+
+if __name__ == "__main__":
     main()