da03 commited on
Commit
02df9f8
1 Parent(s): 0375864

Add application file

Browse files
Files changed (2) hide show
  1. app.py +31 -0
  2. requirements.txt +1 -0
app.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import spaces
2
+ import gradio as gr
3
+ from transformers import AutoTokenizer, AutoModelForCausalLM
4
+
5
+ model_name = 'yuntian-deng/gpt2-small-implicit-cot-multiplication'
6
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
7
+ model = AutoModelForCausalLM.from_pretrained(model_name)
8
+
9
+ def preprocess(num):
10
+ num = num.strip().replace(' ', '')
11
+ reversed_num = ' '.join(num[::-1])
12
+ return reversed_num
13
+
14
+ @spaces.GPU
15
+ def predict_product(num1, num2):
16
+ input_text = f'{preprocess(num1)} * {preprocess(num2)} ='
17
+ inputs = tokenizer(input_text, return_tensors='pt').to('cuda' if torch.cuda.is_available() else 'cpu')
18
+ model.to('cuda' if torch.cuda.is_available() else 'cpu')
19
+ outputs = model.generate(**inputs, max_new_tokens=40)
20
+ prediction = tokenizer.decode(outputs[0], skip_special_tokens=True)
21
+ return prediction.strip()
22
+
23
+ demo = gr.Interface(
24
+ fn=predict_product,
25
+ inputs=[gr.Number(label='First Number (up to 9 digits)'), gr.Number(label='Second Number (up to 9 digits)')],
26
+ outputs='text',
27
+ title='GPT-2 Multiplication Predictor',
28
+ description='Enter two numbers up to 9 digits each and get the predicted product.'
29
+ )
30
+
31
+ demo.launch()
requirements.txt ADDED
@@ -0,0 +1 @@
 
 
1
+ transformers