Update app.py

app.py

@@ -4,68 +4,68 @@ import torch
 import json
 
 # load all models
-# pragformer = transformers.AutoModel.from_pretrained("Pragformer/PragFormer", trust_remote_code=True)
-# pragformer_private = transformers.AutoModel.from_pretrained("Pragformer/PragFormer_private", trust_remote_code=True)
-# pragformer_reduction = transformers.AutoModel.from_pretrained("Pragformer/PragFormer_reduction", trust_remote_code=True)
+pragformer = transformers.AutoModel.from_pretrained("Pragformer/PragFormer", trust_remote_code=True)
+pragformer_private = transformers.AutoModel.from_pretrained("Pragformer/PragFormer_private", trust_remote_code=True)
+pragformer_reduction = transformers.AutoModel.from_pretrained("Pragformer/PragFormer_reduction", trust_remote_code=True)
 
 
 #Event Listeners
 
-# tokenizer = transformers.AutoTokenizer.from_pretrained('NTUYG/DeepSCC-RoBERTa')
+tokenizer = transformers.AutoTokenizer.from_pretrained('NTUYG/DeepSCC-RoBERTa')
 
-# with open('c_data.json', 'r') as f:
-#     data = json.load(f)
+with open('c_data.json', 'r') as f:
+    data = json.load(f)
 
-# def fill_code(code_pth):
-#     return data[code_pth]['pragma'], data[code_pth]['code']
+def fill_code(code_pth):
+    return data[code_pth]['pragma'], data[code_pth]['code']
   
 
-# def predict(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(torch.tensor(tokenized['input_ids']), torch.tensor(tokenized['attention_mask']))
-
-#     y_hat = torch.argmax(pred).item()
-#     return 'With OpenMP' if y_hat==1 else 'Without OpenMP', torch.nn.Softmax(dim=1)(pred).squeeze()[y_hat].item()
-
-
-# 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)
+def predict(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(torch.tensor(tokenized['input_ids']), torch.tensor(tokenized['attention_mask']))
+
+    y_hat = torch.argmax(pred).item()
+    return 'With OpenMP' if y_hat==1 else 'Without OpenMP', torch.nn.Softmax(dim=1)(pred).squeeze()[y_hat].item()
+
+
+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)
 
 
 # Define GUI
@@ -118,10 +118,10 @@ with gr.Blocks() as pragformer_gui:
                 private = gr.Textbox(label="Private", visible=False)
                 reduction = gr.Textbox(label="Reduction", visible=False)
 
-    # submit_btn.click(fn=predict, inputs=code_in, outputs=[label_out, confidence_out])
-    # submit_btn.click(fn=is_private, inputs=code_in, outputs=private)
-    # submit_btn.click(fn=is_reduction, inputs=code_in, outputs=reduction)
-    # sample_btn.click(fn=fill_code, inputs=drop, outputs=[pragma, code_in])
+    submit_btn.click(fn=predict, inputs=code_in, outputs=[label_out, confidence_out])
+    submit_btn.click(fn=is_private, inputs=code_in, outputs=private)
+    submit_btn.click(fn=is_reduction, inputs=code_in, outputs=reduction)
+    sample_btn.click(fn=fill_code, inputs=drop, outputs=[pragma, code_in])
 
 
 pragformer_gui.launch()