Update app.py

app.py

@@ -1,80 +1,38 @@
 import gradio as gr
 from transformers import RobertaTokenizer, RobertaForSequenceClassification
 import torch
-import psutil
-import threading
-from queue import Queue
 
 # Load the model and tokenizer from the specified directory
 model_path = './finetuned_roberta'
 tokenizer = RobertaTokenizer.from_pretrained(model_path)
 model = RobertaForSequenceClassification.from_pretrained(model_path)
 
-# Initialize a request queue with a maximum of 2 concurrent requests
-request_queue = Queue(maxsize=2)
-
-# Function to get CPU and RAM usage
-def get_system_usage():
-    cpu_usage = psutil.cpu_percent(interval=1)
-    ram_usage = psutil.virtual_memory().percent
-    return f"CPU Usage: {cpu_usage}%", f"RAM Usage: {ram_usage}%"
-
-# Function to get the user's position in the queue
-def get_queue_position():
-    return f"Queue Position: {request_queue.qsize() + 1}"
-
 # Define the prediction function
 def classify_text(text):
-    request_queue.put(1)  # Add request to queue
-    position_in_queue = get_queue_position()
+    # Tokenize the input text
+    inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=128)
     
-    while request_queue.full():
-        position_in_queue = get_queue_position()
-        
-    try:
-        # Tokenize the input text
-        inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=128)
-        
-        # Get the model's prediction
-        with torch.no_grad():
-            outputs = model(**inputs)
-        
-        # Apply softmax to get probabilities
-        probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
-        
-        # Get the probability of the class '1'
-        prob_1 = probabilities[0][1].item()
-        
-        return {"Probability of being 1": prob_1, "Queue Position": position_in_queue}
-    finally:
-        request_queue.get()  # Remove request from queue
-
-# Create the Gradio interface
-with gr.Blocks() as iface:
-    with gr.Row():
-        gr.Markdown("### Text Classification with RoBERTa")
+    # Get the model's prediction
+    with torch.no_grad():
+        outputs = model(**inputs)
     
-    with gr.Row():
-        with gr.Column():
-            input_text = gr.Textbox(lines=2, placeholder="Enter text here...")
-            classify_btn = gr.Button("Classify")
-        with gr.Column():
-            cpu_output = gr.Markdown("")
-            ram_output = gr.Markdown("")
-            queue_output = gr.Markdown("")
+    # Apply softmax to get probabilities
+    probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
     
-    output_json = gr.JSON()
-
-    def update_usage():
-        while True:
-            cpu_usage, ram_usage = get_system_usage()
-            cpu_output.update(cpu_usage)
-            ram_output.update(ram_usage)
-    
-    threading.Thread(target=update_usage, daemon=True).start()
+    # Get the probability of the class '1'
+    prob_1 = probabilities[0][1].item()
     
-    classify_btn.click(classify_text, inputs=input_text, outputs=output_json)
+    return {"Probability of being 1": prob_1}
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=classify_text,
+    inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),
+    outputs="json",
+    title="Text Classification with RoBERTa",
+    description="Enter some text and get the probability of the text being written by AI.",
+)
 
 # Launch the app
 if __name__ == "__main__":
-    iface.launch(share=True)
+    iface.launch(share=True)