import spaces 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) 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_name, 'no': no_cot_model, 'explicit': explicit_cot_model} # Constants MAX_PRODUCT_DIGITS_PER_MODEL = {'implicit': 100, 'no': 100, 'explicit': 900} 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 @spaces.GPU def predict_product(num1, num2): input_text = f'{preprocess(num1)} * {preprocess(num2)} =' 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()] 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_results_per_model = {} for step in range(max(MAX_PRODUCT_DIGITS_PER_MODEL.values())): # Set a maximum limit to prevent infinite loops # Ground Truth ground_truth_results = [] for i in range(step+1): ground_truth_digit = ground_truth_digits_reversed[i] ground_truth_digits.append((ground_truth_digit, None)) ground_truth_digits = ground_truth_digits[::-1] # Predicted for model_name in models: model = models[model_name] if finished_per_model[model_name]: continue if step >= MAX_PRODUCT_DIGITS_PER_MODE[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: print ('berak') break 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_results = [] is_correct_sofar = True for i in range(len(predicted_digits_reversed)): predicted_digit = predicted_digits_reversed[i] ground_truth_digit = ground_truth_digits_reversed[i] if i >= len(ground_truth_digits_reversed): if predicted_digit == '0' and is_correct_sofar: is_correct_digit = True else: is_correct_digit = False else: 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: predicted_results.append((predicted_digit, "correct")) else: predicted_results.append((predicted_digit, "wrong")) predicted_results = predicted_results[::-1] predicted_results_per_model[model_name] = predicted_results predicted_results_implicit_cot = predicted_results_per_model['implicit'] predicted_results_nocot = predicted_results_per_model['no'] predicted_results_explicit_cot = predicted_results_per_model['explicit'] yield ground_truth_digits_digits, predicted_results_implicit, predicted_results_nocot, predicted_results_explicit_cot demo = gr.Interface( fn=predict_product, inputs=[ gr.Textbox(label='First Number (up to 12 digits)', value='123456789'), gr.Textbox(label='Second Number (up to 12 digits)', value='987654321'), ], color_map = {"correct": "green", "wrong": "red"} outputs=[ gr.HighlightedText(label='Ground Truth Product', combine_adjacent=False, show_legend=False, color_map=color_map), gr.HighlightedText(label='Implicit CoT Predicted Product', combine_adjacent=False, show_legend=False, color_map=color_map, show_inline_category=False), gr.HighlightedText(label='No CoT Predicted Product', combine_adjacent=False, show_legend=False, color_map=color_map, show_inline_category=False), gr.HighlightedText(label='Explicit CoT Predicted Product', combine_adjacent=False, show_legend=False, color_map=color_map, show_inline_category=False), ], title='GPT2 Direct Multiplication Calculator (Without Using Chain-of-Thought)', description='This demo uses GPT2 to directly predict the product of two numbers without using any intermediate reasoning steps. The GPT2 model has been fine-tuned to internalize chain-of-thought reasoning within its hidden states, following our stepwise internalization approach detailed in the paper linked at the bottom of this page.', article=""" - [Paper: From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step](https://arxiv.org/pdf/2405.14838) - [Code Repository](https://github.com/da03/Internalize_CoT_Step_by_Step) - [Tweet Announcement](https://twitter.com/yuntiandeng/status/1795854740879774036) """, clear_btn=None, submit_btn="Multiply!", live=False ) demo.launch()