Update app.py

app.py

@@ -10,7 +10,6 @@ pragformer_reduction = transformers.AutoModel.from_pretrained("Pragformer/PragFo
 
 
 #Event Listeners
-
 tokenizer = transformers.AutoTokenizer.from_pretrained('NTUYG/DeepSCC-RoBERTa')
 
 with open('c_data.json', 'r') as f:
@@ -35,37 +34,43 @@ def predict(code_txt):
 
 
 def is_private(code_txt):
-    code = code_txt.lstrip().rstrip()
-    tokenized = tokenizer.batch_encode_plus(
-                [code],
-                max_length = 150,
-                pad_to_max_length = True,
-                truncation = True
-            )
-    pred = pragformer_private(torch.tensor(tokenized['input_ids']), torch.tensor(tokenized['attention_mask']))
-
-    y_hat = torch.argmax(pred).item()
-    if y_hat == 0:
-        return gr.update(visible=False)
-    else:
-        return gr.update(value=f"Confidence: {torch.nn.Softmax(dim=1)(pred).squeeze()[y_hat].item()}", visible=True)
-
-
-def is_reduction(code_txt):
-    code = code_txt.lstrip().rstrip()
-    tokenized = tokenizer.batch_encode_plus(
-                [code],
-                max_length = 150,
-                pad_to_max_length = True,
-                truncation = True
-            )
-    pred = pragformer_reduction(torch.tensor(tokenized['input_ids']), torch.tensor(tokenized['attention_mask']))
-
-    y_hat = torch.argmax(pred).item()
-    if y_hat == 0:
-        return gr.update(visible=False)
-    else:
-        return gr.update(value=f"Confidence: {torch.nn.Softmax(dim=1)(pred).squeeze()[y_hat].item()}", visible=True)
+  if predict(code_txt)[0] == 'Without OpenMP':
+      return gr.update(visible=False)
+
+  code = code_txt.lstrip().rstrip()
+  tokenized = tokenizer.batch_encode_plus(
+            [code],
+            max_length = 150,
+            pad_to_max_length = True,
+            truncation = True
+        )
+  pred = pragformer_private(torch.tensor(tokenized['input_ids']), torch.tensor(tokenized['attention_mask']))
+
+  y_hat = torch.argmax(pred).item()
+  if y_hat == 0:
+      return gr.update(visible=False)
+  else:
+      return gr.update(value=f"Confidence: {torch.nn.Softmax(dim=1)(pred).squeeze()[y_hat].item()}", visible=True)
+
+
+def is_reduction(code_txt, label):
+  if predict(code_txt)[0] == 'Without OpenMP':
+      return gr.update(visible=False)
+      
+  code = code_txt.lstrip().rstrip()
+  tokenized = tokenizer.batch_encode_plus(
+            [code],
+            max_length = 150,
+            pad_to_max_length = True,
+            truncation = True
+        )
+  pred = pragformer_reduction(torch.tensor(tokenized['input_ids']), torch.tensor(tokenized['attention_mask']))
+
+  y_hat = torch.argmax(pred).item()
+  if y_hat == 0:
+      return gr.update(visible=False)
+  else:
+      return gr.update(value=f"Confidence: {torch.nn.Softmax(dim=1)(pred).squeeze()[y_hat].item()}", visible=True)
 
 
 # Define GUI
@@ -102,7 +107,7 @@ with gr.Blocks() as pragformer_gui:
             gr.Markdown("## Input")
             with gr.Row():
                 with gr.Column():
-                    drop = gr.Dropdown(list(data.keys()), label="Random Code Snippet") 
+                    drop = gr.Dropdown(list(data.keys()), label="Random Code Snippet", value="LLNL/AutoParBench/benchmarks/Autopar/NPB3.0-omp-c/BT/bt/129")
                     sample_btn = gr.Button("Sample")
                 
                 pragma =  gr.Textbox(label="Pragma")