Update app.py

app.py

@@ -1,6 +1,5 @@
 import logging
 import gradio as gr
-from queue import Queue
 import time
 from prometheus_client import start_http_server, Counter, Histogram, Gauge
 import threading
@@ -8,12 +7,15 @@ import psutil
 import random
 from transformers import pipeline
 from sklearn.metrics import precision_score, recall_score, f1_score
-import json
 import requests
+from datasets import load_dataset
 
 # Load the model
 ner_pipeline = pipeline("ner", model="Sevixdd/roberta-base-finetuned-ner")
 
+# Load the dataset
+dataset = load_dataset("surrey-nlp/PLOD-filtered")
+
 # --- Prometheus Metrics Setup ---
 REQUEST_COUNT = Counter('gradio_request_count', 'Total number of requests')
 REQUEST_LATENCY = Histogram('gradio_request_latency_seconds', 'Request latency in seconds')
@@ -30,29 +32,35 @@ logging.basicConfig(filename="chat_log.txt", level=logging.DEBUG, format='%(asct
 chat_queue = Queue()  # Define chat_queue globally
 
 # --- Chat Function with Monitoring ---
-def chat_function(message, user_ner_tags, ground_truth):
+def chat_function(index):
     logging.debug("Starting chat_function")
     with REQUEST_LATENCY.time():
         REQUEST_COUNT.inc()
         try:
-            chat_queue.put(message)
-            logging.info(f"Received message from user: {message}")
+            chat_queue.put(index)
+            logging.info(f"Received index from user: {index}")
+
+            # Get the example from the dataset
+            example = dataset['train'][int(index)]
+            tokens = example['tokens']
+            ground_truth_labels = example['ner_tags']
+
+            logging.info(f"Tokens: {tokens}")
+            logging.info(f"Ground Truth Labels: {ground_truth_labels}")
 
-            ner_results = ner_pipeline(message)
+            # Predict using the model
+            ner_results = ner_pipeline(" ".join(tokens))
             logging.debug(f"NER results: {ner_results}")
 
             detailed_response = []
             model_predicted_labels = []
-            user_predicted_labels = []
             for result in ner_results:
                 token = result['word']
                 score = result['score']
                 entity = result['entity']
-                start = result['start']
-                end = result['end']
                 label_id = int(entity.split('_')[-1])  # Extract numeric label from entity
                 model_predicted_labels.append(label_id)
-                detailed_response.append(f"Token: {token}, Entity: {entity}, Score: {score:.4f}, Start: {start}, End: {end}")
+                detailed_response.append(f"Token: {token}, Entity: {entity}, Score: {score:.4f}")
 
             response = "\n".join(detailed_response)
             logging.info(f"Generated response: {response}")
@@ -62,21 +70,12 @@ def chat_function(message, user_ner_tags, ground_truth):
 
             time.sleep(random.uniform(0.5, 2.5))  # Simulate processing time
 
-            # Compare user's input tags to the model's output
-            try:
-                user_ner_results = json.loads(user_ner_tags)
-                if not isinstance(user_ner_results, list):
-                    raise ValueError("Invalid format for user NER tags. Please provide a JSON list of dictionaries.")
-                for result in user_ner_results:
-                    entity = result['entity']
-                    label_id = int(entity.split('_')[-1])  # Extract numeric label from entity
-                    user_predicted_labels.append(label_id)
-            except json.JSONDecodeError:
-                return "Invalid JSON format for user NER tags. Please provide a valid JSON array."
-
-            precision = precision_score(user_predicted_labels, model_predicted_labels, average='weighted', zero_division=0)
-            recall = recall_score(user_predicted_labels, model_predicted_labels, average='weighted', zero_division=0)
-            f1 = f1_score(user_predicted_labels, model_predicted_labels, average='weighted', zero_division=0)
+            # Ensure the model and ground truth labels are the same length for comparison
+            model_predicted_labels = model_predicted_labels[:len(ground_truth_labels)]
+
+            precision = precision_score(ground_truth_labels, model_predicted_labels, average='weighted', zero_division=0)
+            recall = recall_score(ground_truth_labels, model_predicted_labels, average='weighted', zero_division=0)
+            f1 = f1_score(ground_truth_labels, model_predicted_labels, average='weighted', zero_division=0)
 
             metrics_response = (f"Precision: {precision:.4f}\n"
                                 f"Recall: {recall:.4f}\n"
@@ -93,10 +92,10 @@ def chat_function(message, user_ner_tags, ground_truth):
             return f"An error occurred. Please try again. Error: {e}"
 
 # Function to simulate stress test
-def stress_test(num_requests, message, delay):
+def stress_test(num_requests, index, delay):
     def send_chat_message():
         response = requests.post("http://127.0.0.1:7860/api/predict/", json={
-            "data": [message],
+            "data": [index],
             "fn_index": 0  # This might need to be updated based on your Gradio app's function index
         })
         logging.debug(response.json())
@@ -121,10 +120,9 @@ body {
 """, title="PLOD Filtered with Monitoring") as demo:  # Load CSS for background image
     with gr.Tab("Chat"):
         gr.Markdown("## Chat with the Bot")
-        message_input = gr.Textbox(label="Enter your sentence:", lines=2)
-        user_ner_tags_input = gr.Textbox(label="Enter your NER tags (JSON format):", lines=5)
-        output = gr.Textbox(label="Response", lines=10)
-        chat_interface = gr.Interface(fn=chat_function, inputs=[message_input, user_ner_tags_input, gr.Textbox(lines=5)], outputs=output)
+        index_input = gr.Textbox(label="Enter dataset index:", lines=1)
+        output = gr.Textbox(label="Response", lines=20)
+        chat_interface = gr.Interface(fn=chat_function, inputs=[index_input], outputs=output)
         chat_interface.render()
 
     with gr.Tab("Model Parameters"):
@@ -143,20 +141,20 @@ body {
 
     with gr.Tab("Stress Testing"):
         num_requests_input = gr.Number(label="Number of Requests", value=10)
-        message_input_stress = gr.Textbox(label="Message", value="Hello bot!")
+        index_input_stress = gr.Textbox(label="Dataset Index", value="2")
         delay_input = gr.Number(label="Delay Between Requests (seconds)", value=0.1)
         stress_test_button = gr.Button("Start Stress Test")
         stress_test_status = gr.Textbox(label="Stress Test Status", lines=5, interactive=False)
 
-        def run_stress_test(num_requests, message, delay):
+        def run_stress_test(num_requests, index, delay):
             stress_test_status.value = "Stress test started..."
             try:
-                stress_test(num_requests, message, delay)
+                stress_test(num_requests, index, delay)
                 stress_test_status.value = "Stress test completed."
             except Exception as e:
                 stress_test_status.value = f"Stress test failed: {e}"
 
-        stress_test_button.click(run_stress_test, [num_requests_input, message_input_stress, delay_input], stress_test_status)
+        stress_test_button.click(run_stress_test, [num_requests_input, index_input_stress, delay_input], stress_test_status)
 
     # --- Update Functions ---
     def update_metrics(request_count_display, avg_latency_display):