import streamlit as st import torch from transformers import T5ForConditionalGeneration, RobertaTokenizer import re import ast # Load the fine-tuned model and tokenizer model_repo_path = 'sabssag/Latex_to_Python_CodeT5-base' model = T5ForConditionalGeneration.from_pretrained(model_repo_path, torch_dtype=torch.float16) tokenizer = RobertaTokenizer.from_pretrained(model_repo_path) model.eval() # Fix unmatched brackets def fix_unmatched_brackets(code): """ Fix unmatched brackets in the code by ensuring that all opening brackets have corresponding closing brackets, and ensure that newline characters are handled correctly when adding missing brackets. """ open_brackets = {'(': 0, '[': 0, '{': 0} close_brackets = {')': 0, ']': 0, '}': 0} bracket_pairs = {'(': ')', '[': ']', '{': '}'} stack = [] new_code = "" # Iterate through the code to track unmatched brackets and their positions for i, char in enumerate(code): if char in open_brackets: stack.append(char) open_brackets[char] += 1 elif char in close_brackets: if stack and bracket_pairs[stack[-1]] == char: stack.pop() # Matching bracket else: # Unmatched closing bracket found, but we need to check if it's valid if stack: # If we have an unmatched opening bracket, fix it by adding the correct closing bracket new_code += bracket_pairs[stack.pop()] else: # If no matching opening bracket, just skip adding the closing bracket continue new_code += char # Append missing closing brackets at the end while stack: last_char = new_code[-1] # If the last character is a newline, remove it before appending the closing bracket if last_char == '\n': new_code = new_code[:-1] new_code += bracket_pairs[stack.pop()] return new_code # Validate and correct bracket balance def validate_bracket_balance(code): """ Validates if brackets are balanced and fixes common issues. """ stack = [] bracket_map = {')': '(', ']': '[', '}': '{'} for i, char in enumerate(code): if char in bracket_map.values(): stack.append(char) elif char in bracket_map: if stack and stack[-1] == bracket_map[char]: stack.pop() else: code = code[:i] + '#' + code[i+1:] # Comment out the misaligned closing bracket break while stack: code += { '(': ')', '[': ']', '{': '}' }[stack.pop()] return code # Add missing imports based on used functions def add_missing_imports(code): """ Detect missing sympy or numpy imports based on used functions in the code. Also fixes incorrect import statements like `from sympy import, pi`. """ sympy_funcs = { "cot", "sqrt", "pi", "sin", "cos", "tan", "log", "Abs", "exp", "factorial", "csc", "sec", "asin", "acos", "atan", "Eq", "symbols", "Function", "Derivative" } # Detect function calls and existing imports function_pattern = r'\b([a-zA-Z_][a-zA-Z0-9_]*)\b' used_functions = set(re.findall(function_pattern, code)) # Match 'from sympy import' statements existing_imports = re.findall(r'from sympy import ([a-zA-Z_, ]+)', code) # Flatten the existing imports set by splitting any comma-separated imports existing_imports_set = {imp.strip() for ex_imp in existing_imports for imp in ex_imp.split(',')} # Find which sympy functions are required but not yet imported required_imports = used_functions.intersection(sympy_funcs) - existing_imports_set # If there are required imports, we will just add them on top of the existing imports if required_imports: # Consolidate all imports into one line, without adding duplicate imports import_statement = f"from sympy import {', '.join(sorted(existing_imports_set | required_imports))}\n" # Remove the current sympy imports with a consolidated import statement code = re.sub(r'from sympy import [a-zA-Z_, ]+\n', '', code) code = import_statement + code # Fully remove incorrect import statements (like `from sympy import, pi`) code = re.sub(r'from sympy import,\s*.*\n', '', code) # Add numpy import if necessary if "np." in code and "import numpy as np" not in code: code = "import numpy as np\n" + code return code # Enhanced removal of evalf() calls, handling malformed cases def remove_evalf(code): """ Remove all occurrences of .evalf() from the code, including cases where it's misplaced or malformed. """ # Remove evalf calls in a more comprehensive way code = re.sub(r'\.evalf\(\)', '', code) # Regular evalf calls code = re.sub(r'\*evalf\(\)', '', code) # Cases like `*evalf()` # Ensure parentheses remain balanced even after removing evalf() code = fix_unmatched_brackets(code) return code def handle_sum_errors(code): """ Detects and fixes cases where `sum()` is applied to non-iterable objects. """ # Regex to detect invalid use of sum invalid_sum_pattern = r'sum\(([^()]+)\)' # Replace invalid sum usage with the content inside the sum (since it's non-iterable) code = re.sub(invalid_sum_pattern, r'\1', code) return code def complete_try_catch_block(code): """ Ensure that the try block in the code is followed by a valid except block. If missing, a generic except block will be added. """ # Check if there's a 'try' block without an 'except' block if 'try:' in code and 'except' not in code: # Add a generic except block to catch any exceptions code = re.sub(r'try:', r'try:\n pass\n except Exception as e:\n print(f"Error: {e}")', code) return code import re def remove_extra_variables_from_function(code): """ Remove extra variables from the function definition list of arguments that are not used in the function body. """ # Find the function definition match = re.search(r'def\s+([a-zA-Z_][a-zA-Z0-9_]*)\((.*?)\):', code) if match: func_name = match.group(1) arg_list = match.group(2).split(',') arg_list = [arg.strip() for arg in arg_list] # Clean up spaces # Get the body of the function (everything after the definition) func_body = code.split(':', 1)[1] # Find which variables are actually used in the function body used_vars = set(re.findall(r'\b([a-zA-Z_][a-zA-Z0-9_]*)\b', func_body)) # Filter out only the arguments that are actually used in the function body filtered_args = [arg for arg in arg_list if arg in used_vars] # Reconstruct the function definition with only the used arguments new_func_def = f"def {func_name}({', '.join(filtered_args)}):" # Replace the old function definition with the new one code = re.sub(r'def\s+[a-zA-Z_][a-zA-Z0-9_]*\s*\(.*?\):', new_func_def, code) return code # Post-process the generated code def post_process_code(code): code = fix_unmatched_brackets(code) code = validate_bracket_balance(code) code = add_missing_imports(code) code = remove_evalf(code) code = handle_sum_errors(code) code = complete_try_catch_block(code) code = remove_extra_variables_from_function(code) return code # Generate the final code from LaTeX def generate_code(latex_expression, max_length=512): inputs = tokenizer(f"Latex Expression: {latex_expression} Solution:", return_tensors="pt") outputs = model.generate(**inputs, max_length=max_length) generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True) post_processed_code = post_process_code(generated_code) return post_processed_code # Streamlit app layout st.title("LaTeX to Python Code Generator") # Define session state keys if 'latex_expr' not in st.session_state: st.session_state.latex_expr = "" # User input for LaTeX expression latex_input = st.text_area("Enter the LaTeX Expression", value=st.session_state.latex_expr, height=150) # Update session state with the new LaTeX expression if st.button("Generate Code"): if latex_input: st.session_state.latex_expr = latex_input with st.spinner("Generating Python Code..."): try: # Correct function name here generated_code = generate_code(latex_expression=st.session_state.latex_expr) # Display the generated code st.subheader("Generated Python Code") st.code(generated_code, language='python') except Exception as e: st.error(f"Error during code generation: {e}") else: st.warning("Please enter a LaTeX expression to generate Python code.")