import streamlit as st
import torch
import numpy as np
import io
import sys

# Function to execute the input code and capture print statements
def execute_code(code):
    # Redirect stdout to capture print statements
    old_stdout = sys.stdout
    sys.stdout = mystdout = io.StringIO()
    
    global_vars = {"torch": torch, "np": np}
    local_vars = {}
    try:
        exec(code, global_vars, local_vars)
        output = mystdout.getvalue()
    except Exception as e:
        output = str(e)
    finally:
        # Reset redirect.
        sys.stdout = old_stdout
    
    return output, local_vars

st.title('PyTorch Code Runner')

# Text area for inputting the PyTorch code
code_input = st.text_area("Enter your PyTorch code here", height=300)

# Button to execute the code
if st.button("Run Code"):
    # Prepend the import statement
    code_to_run = "import torch\nimport numpy as np\n" + code_input
    
    # Execute the code and capture the output
    output, variables = execute_code(code_to_run)
    
    # Display the output
    st.subheader('Output')
    st.text(output)