app.py

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.")