Basic interface

app.py

@@ -1,7 +1,42 @@
 import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-def test(name):
-    return "test"
+tokenizer = AutoTokenizer.from_pretrained("georgiyozhegov/calculator-8m")
+model = AutoModelForCausalLM.from_pretrained("georgiyozhegov/calculator-8m")
+
+def solve(problem):
+    prompt = f"find {problem}\nstep"
+    inputs = tokenizer(prompt, return_tensors="pt", return_token_type_ids=False)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+            max_length=32,
+            do_sample=True,
+            top_k=50,
+            top_p=0.98
+        )
+
+    count = 0
+    for index, token in enumerate(outputs[0]):
+        if token == 6: count += 1
+        if count >= 2: break
+
+    output = tokenizer.decode(outputs[0][:index])
+    return output
+
+examples = [
+    ["2 + 3"],
+    ["10 / 0.5"],
+]
+
+demo = gr.Interface(
+    fn=solve,
+    inputs=gr.Textbox(lines=5, label="Problem"),
+    outputs=gr.Textbox(label="Solution"),
+    examples=examples,
+)
 
-demo = gr.Interface(fn=test, inputs="text", outputs="text")
 demo.launch()