Create app.py

app.py

@@ -0,0 +1,45 @@
+from transformers import AutoModelForCausalLM
+import torch
+import gradio as gr
+
+model = AutoModelForCausalLM.from_pretrained("models/addition_model")
+
+class NumberTokenizer:
+  def __init__(self, numbers_qty=10):
+    vocab = ['+', '=', '-1', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
+    self.numbers_qty = numbers_qty
+    self.pad_token = '-1'
+    self.encoder = {str(v):i for i,v in enumerate(vocab)}
+    self.decoder = {i:str(v) for i,v in enumerate(vocab)}
+    self.pad_token_id = self.encoder[self.pad_token]
+
+  def decode(self, token_ids):
+    return ' '.join(self.decoder[t] for t in token_ids)
+
+  def __call__(self, text):
+    return [self.encoder[t] for t in text.split()]
+
+tokenizer = NumberTokenizer(13)
+
+def generate_solution(input, solution_length=6, model=model):
+  model.eval()
+  input = torch.tensor(tokenizer(input))
+  input = input
+  solution = []
+  for i in range(solution_length):
+    output = model(input)
+    predicted = output.logits[-1].argmax()
+    input = torch.cat((input, predicted.unsqueeze(0)), dim=0)
+    solution.append(predicted.cpu().item())
+  return tokenizer.decode(solution)
+
+def solve(input):
+    return generate_solution(input, solution_length=2)
+
+demo = gr.Interface(fn=solve, inputs=[gr.Textbox(label="Addition exercise", lines=1, info="The input must be of the form '1 + 2 =', with a single space between each character, and only single-digit numbers are allowed.")],
+                    outputs=[gr.Textbox(label="Result", lines=1)],
+                    title="Simple addition with a GPT-like model",
+                    description="Perform addition of two single-digit numbers using a GPT-like model trained on a small dataset.",
+                    examples=["1 + 2 =", "5 + 7 ="])
+    
+demo.launch()