1. 1
2. 2
3. 3
4. 4
5. 5
6. 6
7. 7
8. 8
9. 9
10. 10
11. 11
12. 12
13. 13
14. 14
15. 15
16. 16
17. 17
18. 18
19. 19
20. 20
21. 21
22. 22
23. 23
24. 24
25. 25
26. 26
27. 27
28. 28
29. 29
30. 30
31. 31
32. 32
33. 33
34. 34
35. 35
36. 36
37. 37
38. 38
39. 39
40. 40
41. 41
42. 42
43. 43
44. 44
45. 45
46. 46
47. 47
48. 48
49. 49
50. 50
51. 51
52. 52
53. 53
54. 54
55. 55
56. 56
57. 57
58. 58
59. 59
60. 60
61. 61
62. 62
63. 63
64. 64
65. 65
66. 66
67. 67
68. 68
69. 69
70. 70
71. 71

from time import time
from IPython.core.magic import register_line_magic
from requests import post


@register_line_magic
def nl(line, testing=False):
    """
    Natural language input for Sage using Ollama

    Do not use with untrusted input!

    Example usage:
    %nl Invert the 4x4 matrix with 0s on the diagonal and 1s off diagonal
    """
    r = post(
        "http://localhost:11434/api/generate",
        json={
            "model": "gemma3:27b",
            "system": "Translate this request directly into SageMath code. Be very concise and don't output Markdown code blocks or anything extraneous, only output the raw SageMath code. Make sure it's valid SageMath syntax!",
            "prompt": line,
            "stream": False,
        },
    )
    code = r.json()["response"]
    print(code)
    if testing or input("Run? (y/N) ") == "y":
        shell = get_ipython()
        res = shell.run_cell(code).result
        if testing:
            return res


nl_tests = [
    ("Determinant of the 4x4 matrix with 0s on the diagonal and 1s off diagonal", -3),
    ("Solve x^2 = 2", [x == -sqrt(2), x == sqrt(2)]),
    ("Is 42069 prime", False),
    (
        r"Take \begin{bmatrix}\frac{9}{10} & \frac{1}{10}\\\frac{3}{10} & \frac{7}{10}\end{bmatrix} to the power of 10 and get the answer as a decimal",
        matrix(
            [
                [0.751511654400000, 0.248488345600000],
                [0.745465036800000, 0.254534963200000],
            ]
        ),
    ),
    (
        r"\frac{1^3e^{-1}}{3!}\frac{(1/2)^0e^{-1/2}{0!}\frac{1/2)^0e^{-1/2}}{0!}",
        e ^ (-2) / 6,
    ),
    ("Find the minimum of x^x on 0 to 10", (0.6922006275553464, 0.3678794331853406)),
    ("sum_{n=0}^oo 1/n! x^n", e ^ x),
    ("int 0 to x 2e^(-2y) dy", 1 - e ^ (-2 * x)),
    ("Number of partitions of 69", 3554345),
    ("solve dy/dx+y-1 = 0 where y(10) = 2", (e ^ 10 + e ^ x) * e ^ (-x)),
    (r"e^{\log(\sqrt2)-3i\pi/4} rectangular form", -1 - i),
    ("Answer to the ultimate question of life the universe and everything", 42),
]


def run_nl_tests():
    start_time = time()
    passes = 0
    for line, ans in nl_tests:
        print(line)
        for _ in range(5):
            if ans == nl(line, True):
                passes += 1
    total = len(nl_tests) * 5.0
    print(passes / total)
    print((time() - start_time) / total)