Update app.py

app.py

@@ -6,7 +6,7 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from prettytable import PrettyTable
 import streamlit as st
 
-st.title('Code Generation on the CoNaLa Dataset')
+#st.title('Code Generation on the CoNaLa Dataset')
 
 class CodeGenerator:
     def __init__(self):
@@ -110,6 +110,66 @@ class CodeGenerator:
 
 
 
+
+
+
+import autopep8
+import black
+import isort
+import pylint.lint
+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
+        with open(filename, 'r') as f:
+            self.code = f.read()
+
+    def format(self):
+        try:
+            # Use isort to sort and organize the imports
+            formatted_code = isort.code(self.code)
+
+            # Use black to format the code
+            formatted_code = black.format_str(formatted_code, mode=black.Mode())
+
+            # Use autoimport to add a missing import statement
+            formatted_code = autoimport.fix_code(formatted_code)
+
+            # Use autopep8 to fix any remaining issues
+            formatted_code = autopep8.fix_code(formatted_code)
+
+        except RuntimeError as error:
+            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
+
+        return formatted_code
+
+
+    def save(self, filename):
+        # Save the fixed code to a file
+        with open(filename, 'w') as f:
+            f.write(self.code)
+
+
+
+
+
 code_generator = CodeGenerator()
 
 
@@ -130,6 +190,15 @@ def main():
         st.write(error_message)
 
 
+        st.subheader('Error Correction')
+        formatter = PythonCodeFormatter(output_code)
+        formatted_code = formatter.format()
+        st.write('Code after correction:')
+        st.write(formatted_code)
+
+        
+
+
 if __name__ == '__main__':
     
     main()