Update app.py

majority with steps 1

app.py

@@ -2,10 +2,11 @@ import gradio as gr
 import ctranslate2
 from transformers import AutoTokenizer
 from huggingface_hub import snapshot_download
-from codeexecutor import postprocess_completion, get_majority_vote
+from codeexecutor import get_majority_vote
+import re
 
 # Define the model and tokenizer loading
-model_prompt = "Solve the following mathematical problem: "
+model_prompt = "Explain and solve the following mathematical problem step by step, showing all work: "
 tokenizer = AutoTokenizer.from_pretrained("AI-MO/NuminaMath-7B-TIR")
 model_path = snapshot_download(repo_id="Makima57/deepseek-math-Numina")
 generator = ctranslate2.Generator(model_path, device="cpu", compute_type="int8")
@@ -15,31 +16,59 @@ iterations = 10
 def get_prediction(question):
     input_text = model_prompt + question
     input_tokens = tokenizer.tokenize(input_text)
-    results = generator.generate_batch([input_tokens])
+    results = generator.generate_batch(
+        [input_tokens],
+        max_length=512,
+        sampling_temperature=0.7,
+        sampling_topk=40,
+    )
     output_tokens = results[0].sequences[0]
     predicted_answer = tokenizer.convert_tokens_to_string(output_tokens)
     return predicted_answer
 
-# Function to perform majority voting and solve the problem with steps
+# Function to parse the prediction to extract the answer and steps
+def parse_prediction(prediction):
+    lines = prediction.strip().split('\n')
+    answer = None
+    steps = []
+    for line in lines:
+        # Check for "Answer:" or "answer:"
+        match = re.match(r'^\s*(?:Answer|answer)\s*[:=]\s*(.*)', line)
+        if match:
+            answer = match.group(1).strip()
+        else:
+            steps.append(line)
+    if answer is None:
+        # If no "Answer:" found, assume last line is the answer
+        answer = lines[-1].strip()
+        steps = lines[:-1]
+    steps_text = '\n'.join(steps).strip()
+    return answer, steps_text
+
+# Function to perform majority voting and get steps
 def majority_vote_with_steps(question, num_iterations=10):
     all_predictions = []
-    all_answer = []
-    steps_to_solve = []
-    
+    all_answers = []
+    steps_list = []
+
     for _ in range(num_iterations):
         prediction = get_prediction(question)
-        # Process prediction to get steps and answer
-        answer, success = postprocess_completion(prediction, True, True)
+        answer, steps = parse_prediction(prediction)
         all_predictions.append(prediction)
-        all_answer.append(answer)
-        if success:
-            steps_to_solve.append(answer)  # Add the steps if code executes successfully
-    
-    majority_voted_ans = get_majority_vote(all_answer)
-    
-    # If steps to solve exist, return them, else fallback to "No steps found"
-    steps_solution = steps_to_solve[0] if steps_to_solve else "No steps found"
-    
+        all_answers.append(answer)
+        steps_list.append(steps)
+
+    # Get the majority voted answer
+    majority_voted_ans = get_majority_vote(all_answers)
+
+    # Find the steps corresponding to the majority voted answer
+    for i, ans in enumerate(all_answers):
+        if ans == majority_voted_ans:
+            steps_solution = steps_list[i]
+            break
+    else:
+        steps_solution = "No steps found"
+
     return majority_voted_ans, steps_solution
 
 # Gradio interface for user input and output
@@ -52,7 +81,7 @@ def gradio_interface(question, correct_answer):
         "Correct Solution": correct_answer
     }
 
-# Custom CSS for enhanced design
+# Custom CSS for enhanced design (unchanged)
 custom_css = """
     body {
         background-color: #fafafa;
@@ -94,7 +123,7 @@ custom_css = """
     #math_question, #correct_answer {
         font-size: 20px;
         font-family: 'Poppins', sans-serif;
-        font-weight: 500px;  /* Apply bold */
+        font-weight: 500px;
         color: #007acc;
         margin-bottom: 5px;
         display: inline-block;