app.py

import gradio as gr
import torch


EXAMPLE_MD = """
```python
import torch 

t1 = torch.arange({n1}).view({dim1})

t2 = torch.arange({n2}).view({dim2})

(t1 @ t2).shape = {out_shape} 

```

"""

matrix_loop = """```python
out = 0
for i, j in zip(t1, t2):
    out += i * j
```
"""

def generate_example(dim1: list, dim2: list):
    n1 = 1
    n2 = 1
    for i in dim1:
        n1 *= i
    for i in dim2:
        n2 *= i

    t1 = torch.arange(n1).view(dim1)
    t2 = torch.arange(n2).view(dim2)
    try:
        out_shape = list((t1 @ t2).shape)
    except RuntimeError:
        out_shape = "error"

    code = EXAMPLE_MD.format(
        n1=str(n1), dim1=str(dim1), n2=str(n2), dim2=str(dim2), out_shape=str(out_shape)
    )

    return dim1, dim2, code


def sanitize_dimention(dim):
    if dim is None:
        gr.Error("one of the dimentions is empty, please fill it")
    if "[" in dim:
        dim = dim.replace("[", "")
    if "]" in dim:
        dim = dim.replace("]", "")
    if "," in dim:
        dim = dim.replace(",", " ").strip()
        out = [int(i.strip()) for i in dim.split()]
    else:
        out = [int(dim.strip())]
    if 0 in out:
        gr.Error(
            "Found the number 0 in one of the dimensions which is not allowed, consider using 1 instead"
        )
    return out


def create_row(dim,is_dim=None,checks=None,version=1):
    out = "| "
    n_dim = len(dim)
    for i in range(n_dim):
        if version == 1 :
            # infered last dims
            if (is_dim ==1 and i == n_dim-2) or (is_dim ==2 and i ==n_dim-1):
                color = "green"
                out += f"<strong style='color: {color}'> {dim[i]} </strong>| "
            # check every normal dimension
            elif (is_dim ==1 and i != n_dim-1) or (is_dim ==2 and i ==n_dim-1):
                color = "green" if checks[i] == "V" else "red"
                out += f"<strong style='color: {color}'> {dim[i]} </strong>| "
            # checks last 2 dims
            elif (is_dim ==1 and i == n_dim-1) or (is_dim ==2 and i ==n_dim-2):
                color = "blue" if checks[i] == "V" else "yellow"
                out += f"<strong style='color: {color}'> {dim[i]} </strong>| "
            # when using this function without checks
            else : 
                out+= f"{dim[i]} | "
        if version == 2 : 
            if (is_dim == 1 and i != n_dim-1) : 
                out += f"<strong style='color: green'> {dim[i]} </strong>| "
            elif i == n_dim-1 :
                color = "blue" if checks[i] == "V" else "yellow"
                out += f"<strong style='color: {color}'> {dim[i]} </strong>| "
            else : 
                out += f"{dim[i]} | "
            
    return out + "\n"


def create_header(n_dim, checks=None):
    checks = ["<!-- -->"] * n_dim if checks is None else checks
    out = "| "
    for i in checks:
        out = out + i + " | "
    out += "\n" + "|---" * n_dim + "|\n"
    return out


def generate_table(dim1, dim2, checks=None,version=1):
    n_dim = len(dim1)
    table = create_header(n_dim, checks)
    # tensor 1
    if not checks : 
        table += create_row(dim1)
    else : 
        table += create_row(dim1,1,checks,version)

    # tensor 2
    if not checks : 
        table += create_row(dim2)
    else : 
        table += create_row(dim2,2,checks,version)
    return table


def alignment_and_fill_with_ones(dim1, dim2):
    n_dim = max(len(dim1), len(dim2))

    if len(dim1) == len(dim2):
        pass
    elif len(dim1) < len(dim2):
        placeholder = [1] * (n_dim - len(dim1))
        placeholder.extend(dim1)
        dim1 = placeholder
    else:
        placeholder = [1] * (n_dim - len(dim2))
        placeholder.extend(dim2)
        dim2 = placeholder
    return dim1, dim2

def check_validity(dim1,dim2):
    out = []
    for i in range(len(dim1)-2):
        if dim1[i] == dim2[i]:
            out.append("V")
        else : 
            out.append("X")
    # final dims
    if dim1[-1] == dim2[-2]:
        out.extend(["V","V"])
    else : 
        out.extend(["X","X"])
    return out


def substitute_ones_with_concat(dim1,dim2,version = 1):
    n = len(dim1)-2 if version ==1 else len(dim1)-1 
    for i in range(n):
        dim1[i] = dim2[i] if dim1[i] == 1 else dim1[i]
        dim2[i] = dim1[i] if dim2[i] == 1 else dim2[i]
    return dim1, dim2

def predict(dim1, dim2):
    dim1 = sanitize_dimention(dim1)
    dim2 = sanitize_dimention(dim2)
    n1 , n2 = len(dim1) , len(dim2)
    dim1, dim2, out = generate_example(dim1, dim2)
    # TODO 
    if n1 >1 and n2 > 1 : 
        # Table 1 
        dim1, dim2 = alignment_and_fill_with_ones(dim1, dim2)
        table1 = generate_table(dim1, dim2)
        # Table 2 
        dim1, dim2 = substitute_ones_with_concat(dim1,dim2)
        table2 = generate_table(dim1, dim2)
        # Table 3 
        checks = check_validity(dim1,dim2)
        table3 = generate_table(dim1,dim2,checks)

        out += "\n# Step1 (alignment and pre_append with ones)\n" + table1
        out += "\n# Step2 (susbtitute columns that have 1 with concat)\nexcept for last 2 dimensions\n" + table2
        out += "\n# Step3 (check if matrix multiplication is valid)\n" 
        out += "* last dimension of dim1 should equal before last dimension of dim2 (blue or yellow colors)\n"
        out += "* all the other dimensions should be equal to one another (green or red colors)\n\n" + table3
        if "X" not in checks : 
            dim1[-1] = dim2[-1]
            out += "\n# Final dimension\n"
            out+="as highlighted in <strong style='color:green'> green </strong> \n\n"
            out+= f"`output.shape = {dim1}`"
    # case single dims
    elif n1 == 1 and n2  == 1 :
        out += "# Single Dimensional Cases\n"
        out += "When both matricies have only single dims they should both have the same number of values in the first dimension\n"
        out += "meaning that `t1.shape == t2.shape`\n"
        out += "the output is a single value, think : \n"
        out += matrix_loop
    else : 
        out += "# One of the dimensions is a single dimension\n"
        out += "In this case we need to assert that the last dimension of `t1` "
        out += "is equal to the last dimension of `t2`\n"
        out += "Once the assertion is valid then we get rid of the last dimension and keep the rest\n"
        out += "# Step 1 (alignment and fill with ones)\n"
        dim1, dim2 = alignment_and_fill_with_ones(dim1, dim2)
        table = generate_table(dim1, dim2)
        out += table
        out += "\n# Step2 (susbtitute columns that have 1 with concat)\n"
        out += ""
        dim1, dim2 = substitute_ones_with_concat(dim1,dim2,2)
        checks = ["V"] * (len(dim1)-1)
        if dim1[-1] == dim2[-1] :
            checks.append("V")
        else : 
            checks.append("X")
        table = generate_table(dim1, dim2,checks,2)
        out+= table
        if "X" not in checks : 
            out += "\n#Final dimension"
            out += "The final dimension is everything colored in <strong style='color:green'> green </strong> \n"
            out += f"\nfinal dimension = `{dim1[:-1]}` "
        
        

    return out


demo = gr.Interface(
    predict,
    inputs=["text", "text"],
    outputs=["markdown"],
    examples=[["9,2,1,3,3", "5,3,7"], ["7,4,2,3", "5,2,7"]],
)

demo.launch(debug=True)