Update app.py

app.py

@@ -8,6 +8,14 @@ import streamlit as st
 
 #st.title('Code Generation on the CoNaLa Dataset')
 
+import subprocess
+import re
+import pandas as pd
+import plotly.express as px
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from prettytable import PrettyTable
+
+
 class CodeGenerator:
     def __init__(self):
         self.tokenizer = AutoTokenizer.from_pretrained("AhmedSSoliman/MarianCG-CoNaLa-Large")
@@ -26,7 +34,7 @@ class CodeGenerator:
     def check_code(self, code):
         with open("temp.py", "w") as f:
             f.write(code)
-        result = subprocess.run(["flake8", "temp.py"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        result = subprocess.run(["flake8", "--count", "temp.py"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         output = result.stdout.decode()
         error = result.stderr.decode()
 
@@ -113,6 +121,7 @@ class CodeGenerator:
 
 
 
+
 import autopep8
 import black
 import isort
@@ -123,7 +132,7 @@ import autoimport
 class PythonCodeFormatter:
     def __init__(self, code):
         self.code = code.replace('▁', ' ').strip()
-        
+
 
     def load_code_from_file(self, filename):
         # Load the code to be fixed
@@ -143,21 +152,17 @@ class PythonCodeFormatter:
 
             # Use autopep8 to fix any remaining issues
             formatted_code = autopep8.fix_code(formatted_code)
+            return formatted_code
 
         except RuntimeError as error:
-            if str(error) == 'Project root not found':
+            if str(error) == 'Project root not found.':
                 return formatted_code
             else:
                 raise  # re-raise the error if it's not the one we're looking for
-        except ValueError as error:
-            if str(error) == 'Cannot parse: 2:0: EOF in multi-line statement':
-                # add )) to the end of the line to fix the error
-                self.code += '))'
-                # retry the format() function with the corrected code
-                formatted_code = self.format()
-            else:
-                raise  # re-raise the error if it's not the one we're looking for
 
+        except ValueError as error:
+            return formatted_code
+            
         return formatted_code