.

app.py

@@ -3,156 +3,42 @@ import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Load models
-implicit_cot_model_name = 'yuntian-deng/gpt2-implicit-cot-multiplication'
-implicit_cot_model = AutoModelForCausalLM.from_pretrained(implicit_cot_model_name)
+# Load the implicit CoT model
+implicit_cot_model_name = 'yuntian-deng/implicit-cot-math-mistral7b'
+implicit_cot_model = AutoModelForCausalLM.from_pretrained(implicit_cot_model_name, torch_dtype=torch.bfloat16)
 tokenizer = AutoTokenizer.from_pretrained(implicit_cot_model_name)
 
-no_cot_model_name = 'yuntian-deng/gpt2-no-cot-multiplication'
-no_cot_model = AutoModelForCausalLM.from_pretrained(no_cot_model_name)
-
-explicit_cot_model_name = 'yuntian-deng/gpt2-explicit-cot-multiplication'
-explicit_cot_model = AutoModelForCausalLM.from_pretrained(explicit_cot_model_name)
-
-models = {'implicit': implicit_cot_model, 'no': no_cot_model, 'explicit': explicit_cot_model}
-
 # Constants
-MAX_PRODUCT_DIGITS_PER_MODEL = {'implicit': 100, 'no': 100, 'explicit': 960}
-
-def preprocess(num):
-    num = str(num).strip().replace(' ', '')
-    reversed_num = ' '.join(num[::-1])
-    return reversed_num
-
-def postprocess(raw_output):
-    prediction = raw_output.replace(' ', '')[::-1]
-    return prediction
+MAX_RESULT_TOKENS = 10
 
 @spaces.GPU
-def predict_product(num1, num2):
-    input_text = f'{preprocess(num1)} * {preprocess(num2)} ='
+def predict_answer(question):
+    input_text = ' '.join(question.split()).strip() + ' ' + tokenizer.eos_token
     inputs = tokenizer(input_text, return_tensors='pt').to('cuda' if torch.cuda.is_available() else 'cpu')
-    [model.to('cuda' if torch.cuda.is_available() else 'cpu') for model in models.values()]
+    implicit_cot_model.to('cuda' if torch.cuda.is_available() else 'cpu')
 
     input_ids = inputs['input_ids']
-    input_len = input_ids.shape[-1]
-    prediction = ""
-    ground_truth_product = ""
-    valid_input = True
-
-    try:
-        num1_int = int(num1)
-        num2_int = int(num2)
-        ground_truth_product = str(num1_int * num2_int)
-        ground_truth_digits_reversed = list(ground_truth_product)[::-1]
-    except ValueError:
-        valid_input = False
-
-    generated_ids_per_model = {model_name: inputs['input_ids'].data.clone() for model_name in models}
-    finished_per_model = {model_name: False for model_name in models}
-    past_key_values_per_model = {model_name: None for model_name in models}
-    predicted_annotations_per_model = {}
-    for step in range(max(MAX_PRODUCT_DIGITS_PER_MODEL.values())):  # Set a maximum limit to prevent infinite loops
-        # Ground Truth
-        if not valid_input:
-            ground_truth_annotations = [('Invalid Input!', None)]
-        else:
-            ground_truth_annotations = [(ground_truth_digit, None) for ground_truth_digit in ground_truth_digits_reversed[:step+1]]
-            ground_truth_annotations = ground_truth_annotations[::-1]
-        # Predicted
-        for model_name in models:
-            model = models[model_name]
-            if finished_per_model[model_name]:
-                continue
-            if step >= MAX_PRODUCT_DIGITS_PER_MODEL[model_name]:
-                continue
-            generation_kwargs = {
-                'input_ids': generated_ids_per_model[model_name],
-                'max_new_tokens': 1,
-                'do_sample': False,
-                'past_key_values': past_key_values_per_model[model_name],
-                'return_dict_in_generate': True,
-                'use_cache': True
-            }
-            if step == 0:
-                del generation_kwargs['past_key_values']
-            outputs = model.generate(**generation_kwargs)
-            generated_ids = outputs.sequences
-            next_token_id = generated_ids[0, -1]
-            #print (next_token_id)
-            
-            if next_token_id.item() == tokenizer.eos_token_id:
-                finished_per_model[model_name] = True
-                if valid_input:
-                    if len([item for item in predicted_annotations_per_model[model_name] if item[1] is not None]) < len(ground_truth_digits_reversed):
-                        predicted_annotations_per_model[model_name].insert(0, ('⠀', 'wrong'))
-                continue
-                
-            generated_ids_per_model[model_name] = generated_ids
-            past_key_values_per_model[model_name] = outputs.past_key_values
-            
-            output_text = tokenizer.decode(generated_ids[0, input_len:], skip_special_tokens=True)
-            predicted_digits_reversed = output_text.strip().split(' ')
-        
-            predicted_annotations = []
-            is_correct_sofar = True
-            if model_name == 'explicit':
-                if '=' not in predicted_digits_reversed:
-                    predicted_annotations = [(predicted_digit, None) for predicted_digit in predicted_digits_reversed]
-                    predicted_digits_reversed = []
-                else:
-                    equal_sign_position = predicted_digits_reversed.index('=')
-                    predicted_annotations = [(predicted_digit, None) for predicted_digit in predicted_digits_reversed[:equal_sign_position+1]]
-                    predicted_digits_reversed = predicted_digits_reversed[equal_sign_position+1:]
+    outputs = implicit_cot_model.generate(input_ids=input_ids,
+            max_new_tokens=MAX_RESULT_TOKENS,
+            do_sample=False)
+    generated_ids = outputs.sequences[0]
 
-            for i in range(len(predicted_digits_reversed)):
-                predicted_digit = predicted_digits_reversed[i]
-                if not valid_input:
-                    is_correct_digit = None
-                elif i >= len(ground_truth_digits_reversed):
-                    if predicted_digit == '0' and is_correct_sofar:
-                        is_correct_digit = True
-                    else:
-                        is_correct_digit = False
-                else:
-                    ground_truth_digit = ground_truth_digits_reversed[i]
-                    if predicted_digit == ground_truth_digit:
-                        is_correct_digit = True
-                    else:
-                        is_correct_digit = False
-                if not is_correct_digit:
-                    is_correct_sofar = False
-                if is_correct_digit is None:
-                    predicted_annotations.append((predicted_digit, None))
-                elif is_correct_digit:
-                    predicted_annotations.append((predicted_digit, "correct"))
-                else:
-                    predicted_annotations.append((predicted_digit, "wrong"))
-            predicted_annotations = predicted_annotations[::-1]
-            predicted_annotations_per_model[model_name] = predicted_annotations
+    prediction = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
 
-        predicted_annotations_implicit_cot = predicted_annotations_per_model['implicit']
-        predicted_annotations_nocot = predicted_annotations_per_model['no']
-        predicted_annotations_explicit_cot = predicted_annotations_per_model['explicit']
-
-        yield ground_truth_annotations, predicted_annotations_implicit_cot, predicted_annotations_nocot, predicted_annotations_explicit_cot
+    return prediction
 
 color_map = {"correct": "green", "wrong": "red"}
 
 demo = gr.Interface(
-    fn=predict_product,
+    fn=predict_answer,
     inputs=[
-        gr.Textbox(label='First Number (up to 15 digits)', value='123456789'),
-        gr.Textbox(label='Second Number (up to 15 digits)', value='987654321'),
+        gr.Textbox(label='Question', value='A set of 7 spoons costs $21. If each spoon would be sold separately, how much would 5 spoons cost?'),
     ],
     outputs=[
-        gr.HighlightedText(label='Ground Truth Product', combine_adjacent=False, show_legend=False, color_map=color_map),
-        gr.HighlightedText(label='Implicit CoT Prediction (Ours)', combine_adjacent=False, show_legend=False, color_map=color_map, show_inline_category=False),
-        gr.HighlightedText(label='No CoT Prediction', combine_adjacent=False, show_legend=False, color_map=color_map, show_inline_category=False),
-        gr.HighlightedText(label='Explicit CoT Steps & Prediction', combine_adjacent=False, show_legend=False, color_map=color_map, show_inline_category=False),
+        gr.HighlightedText(label='Implicit CoT Prediction', combine_adjacent=False, show_legend=False, color_map=color_map, show_inline_category=False),
     ],
-    title='Predicting Multiplication with GPT-2: Implicit vs. Explicit CoT',
-    description='This demo showcases GPT-2\'s ability to directly predict the product of two large numbers without intermediate steps, using our stepwise internalization method. Compare the performance of implicit CoT (our method), no CoT, and explicit CoT. Implicit CoT offers accuracy and speed, while explicit CoT provides detailed reasoning but is slower.',
+    title='Solving Grade School Math Problems with Implicit CoT',
+    description='This demo showcases Mistral-7B\'s ability to solve grade school math problems without producing intermediate steps, using our stepwise internalization method.',
     article="""
     - [Paper 1: Implicit Chain of Thought Reasoning via Knowledge Distillation](https://arxiv.org/pdf/2311.01460)
     - [Paper 2: From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step](https://arxiv.org/pdf/2405.14838)
@@ -160,8 +46,8 @@ demo = gr.Interface(
     - [Tweet Announcement](https://twitter.com/yuntiandeng/status/1795854740879774036)
     """,
     clear_btn=None,
-    submit_btn="Multiply!",
+    submit_btn="Get Answer!",
     live=False,
     concurrency_limit=1
 )
-demo.queue(max_size=20).launch()
+demo.queue(max_size=5).launch()