Update README.md

README.md

@@ -3,51 +3,19 @@ license: apache-2.0
 ---
 
 ```
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 import re
-from transformers import AutoTokenizer
-from vllm import LLM, SamplingParams
-import json
-import argparse
-
-# === Argument Parser Setup ===
-# We still allow specifying the model, GPU, etc., but no input/output files.
-parser = argparse.ArgumentParser(description="Demonstrate autoformalization for a single hardcoded example.")
-parser.add_argument('--model_path', default="/scratch/gpfs/st3812/models/Qwen3-32B", type=str, help="Path to the model.")
-parser.add_argument('--seed', default=42, type=int, help="Random seed for reproducibility.")
-parser.add_argument('--gpu', default=1, type=int, help="Number of GPUs for tensor parallelism.")
-parser.add_argument('--max_model_len', default=8192, type=int, help="Maximum model length for VLLM.")
-
-args = parser.parse_args()
-
-# --- 1. Hardcoded Example Problem ---
-# Instead of reading from a file, we define our single problem here.
-problem_entry = {
-    "name": "mathd_algebra_47",
-    "problem": "Let $f(x) = x^2 + 3x + 4$. Find the sum of the roots of $f(x) - 10 = 0$.",
-    "answer": "-3"
-}
-print("📝 Using hardcoded example problem:")
-print(json.dumps(problem_entry, indent=2))
-print("-" * 30)
-
-
-# --- 2. Load Model and Tokenizer ---
-print(f"🚀 Loading model: {args.model_path}...")
-model = LLM(
-    model=args.model_path,
-    seed=args.seed,
-    trust_remote_code=True,
-    tensor_parallel_size=args.gpu,
-    max_model_len=args.max_model_len
-)
-tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
-print("✅ Model and tokenizer loaded.")
-print("-" * 30)
+torch.manual_seed(30)
+
+model_id = "Goedel-LM/Goedel-Formalizer-V2-8B"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True)
 
 
-# --- 3. Prepare Prompt and Sampling Parameters ---
-problem_name = problem_entry.get("name")
-informal_statement_content = f"{problem_entry.get('problem', '')} The answer is {problem_entry.get('answer', '')}"
+problem_name = "test_problem"
+informal_statement_content = "Let $f(x) = x^2 + 3x + 4$. Find the sum of the roots of $f(x) - 10 = 0$. Show the answer is -3."
+
 
 # Construct the prompt for the model
 user_prompt_content = (
@@ -55,32 +23,24 @@ user_prompt_content = (
     f"Use the following theorem name: {problem_name}\n"
     f"The natural language statement is: \n"
     f"{informal_statement_content}"
+    f"Think before you provide the lean statement."
 )
 
-messages = [{"role": "user", "content": user_prompt_content}]
-formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
 
-# Set sampling parameters for generating one sample
-sampling_params = SamplingParams(
-    temperature=0.7,
-    top_p=0.8,
-    max_tokens=2048, # Max new tokens to generate
-    n=1, # Generate exactly one sample
-    seed=args.seed,
-)
 
-print("💬 Generated Prompt:\n" + formatted_prompt)
-print("-" * 30)
+chat = [
+  {"role": "user", "content": user_prompt_content},
+]
 
+inputs = tokenizer.apply_chat_template(chat, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
 
-# --- 4. Run Generation ---
-print("🧠 Generating formal statement...")
-raw_model_outputs = model.generate([formatted_prompt], sampling_params, use_tqdm=False)
-print("✅ Generation complete.")
-print("-" * 30)
+import time
+start = time.time()
+outputs = model.generate(inputs, max_new_tokens=16384, temperature = 0.9, do_sample = True, top_k=50, top_p=0.95)
 
 
-# --- 5. Process and Print the Result ---
+model_output_text = tokenizer.batch_decode(outputs)[0]
+
 def extract_code(text_input):
     """Extracts the last Lean 4 code block from the model's output."""
     try:
@@ -89,19 +49,12 @@ def extract_code(text_input):
     except Exception:
         return "Error during code extraction."
 
-# Since we generated one output for one prompt, we access the first element.
-model_output_text = raw_model_outputs[0].outputs[0].text.strip()
+
 extracted_code = extract_code(model_output_text)
 
-# Prepare a final dictionary with all relevant information
-final_result = {
-    "problem_name": problem_name,
-    "informal_statement": informal_statement_content,
-    "full_model_output": model_output_text,
-    "extracted_lean_code": extracted_code
-}
-
-# Print the final, structured result
-print("🎉 Final Result:")
-print(json.dumps(final_result, indent=4))
+print(time.time() - start)
+print("output:", model_output_text)
+print("lean4 statement:", extracted_code)
+
+
 ```