Upload 2 files

batch_inference.py

@@ -0,0 +1,304 @@
+"""
+Interactive REPL for testing trained physics problem-solving model.
+"""
+import argparse
+from pathlib import Path
+
+import torch
+import yaml
+
+from qwen2_model import Transformer
+from tokenizer import Tokenizer
+from generation_utils import generate
+from tokenizer_wrapper import decode_token_ids
+
+
+SYSTEM_MESSAGE = (
+    "You are a helpful physics tutor. You first think about the reasoning process "
+    "in your mind and then provide the user with the answer."
+)
+USER_TEMPLATE = (
+    "{question}\n"
+    "Show your reasoning in <think> </think> tags. "
+    "Then provide your final answer in <answer> </answer> tags."
+)
+RESPONSE_PROMPT = "Let me solve this step by step.\n<think>"
+
+
+def load_model_and_tokenizer(config_path, checkpoint_path=None):
+    """Load model and tokenizer from config and checkpoint."""
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f)
+    
+    pretrained_model_path = Path(config["model"]["pretrained_model_path"])
+    device = torch.device(config["model"]["device"])
+    
+    dtype_map = {
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+        "float32": torch.float32,
+    }
+    dtype = dtype_map.get(config["model"]["dtype"], torch.bfloat16)
+    
+    # Load tokenizer
+    tokenizer = Tokenizer(str(pretrained_model_path / "tokenizer.json"))
+    
+    # Load model
+    model = Transformer.from_pretrained(pretrained_model_path, device=device)
+    
+    # Load checkpoint if provided
+    if checkpoint_path:
+        print(f"Loading checkpoint from {checkpoint_path}...")
+        checkpoint = torch.load(checkpoint_path, map_location=device)
+        
+        # Handle different checkpoint formats
+        if isinstance(checkpoint, dict):
+            if "model_state_dict" in checkpoint:
+                # Checkpoint contains model_state_dict, optimizer_state_dict, etc.
+                state_dict = checkpoint["model_state_dict"]
+                print(f"Loaded checkpoint from step {checkpoint.get('step', 'unknown')}")
+            else:
+                # Checkpoint is already a state dict
+                state_dict = checkpoint
+        else:
+            state_dict = checkpoint
+        
+        model.load_state_dict(state_dict)
+        print("Checkpoint loaded successfully!")
+    
+    model.eval()
+    
+    return model, tokenizer, device, dtype, config
+
+def generate_response(model, tokenizer, question, device, dtype, max_gen_len=512, temperature=0.7, top_p=0.9):
+    """Generate a response for a given physics question."""
+    # Format the prompt
+    user_message = USER_TEMPLATE.format(question=question)
+    prefix = tokenizer.encode_chat_with_response_prompt(
+        [
+            {"role": "system", "content": SYSTEM_MESSAGE},
+            {"role": "user", "content": user_message},
+        ],
+        RESPONSE_PROMPT,
+    )
+    
+    # Tokenize
+    tokens = tokenizer.tokenize(prefix)
+    prefix_token_ids = tokens.ids
+    
+    # Generate
+    print("\nGenerating response...")
+    with torch.inference_mode():
+        generated_token_ids, is_finished = generate(
+            model=model,
+            tokenizer=tokenizer,
+            prompt_token_ids=prefix_token_ids,
+            max_gen_len=max_gen_len,
+            temperature=temperature,
+            top_p=top_p,
+            device=device,
+            dtype=dtype,
+        )
+    
+    # Decode
+    generated_text = decode_token_ids(tokenizer, generated_token_ids)
+    
+    return prefix + generated_text, is_finished
+
+
+def extract_answer(text):
+    """Extract the answer from <answer> tags."""
+    import re
+    answer_match = re.search(r"<answer>(.*?)</answer>", text, re.DOTALL)
+    if answer_match:
+        return answer_match.group(1).strip()
+    return None
+
+
+def print_response(full_text):
+    """Pretty print the model's response."""
+    import re
+    
+    # Try to extract think and answer sections
+    think_match = re.search(r"<think>(.*?)</think>", full_text, re.DOTALL)
+    answer_match = re.search(r"<answer>(.*?)</answer>", full_text, re.DOTALL)
+    
+    print("\n" + "="*80)
+    
+    if think_match:
+        print("\n🤔 REASONING:")
+        print("-" * 80)
+        print(think_match.group(1).strip())
+    
+    if answer_match:
+        print("\n✅ ANSWER:")
+        print("-" * 80)
+        print(answer_match.group(1).strip())
+    else:
+        print("\n⚠️  WARNING: No answer tags found in response")
+        print("\nFull response:")
+        print("-" * 80)
+        print(full_text)
+    
+    print("="*80 + "\n")
+
+
+def interactive_mode(model, tokenizer, device, dtype, config):
+    """Run interactive REPL mode."""
+    print("\n" + "="*80)
+    print("Physics Problem Solver - Interactive Mode")
+    print("="*80)
+    print("\nCommands:")
+    print("  - Type your physics question and press Enter")
+    print("  - Type 'quit' or 'exit' to exit")
+    print("  - Type 'config' to change generation parameters")
+    print("  - Type 'example' to see example questions")
+    print("="*80 + "\n")
+    
+    # Default generation parameters
+    max_gen_len = config["training"].get("max_gen_len", 512)
+    temperature = 0.7
+    top_p = 0.9
+    
+    while True:
+        try:
+            user_input = input("\n📝 Enter physics question (or command): ").strip()
+            
+            if not user_input:
+                continue
+            
+            if user_input.lower() in ['quit', 'exit', 'q']:
+                print("\nGoodbye! 👋")
+                break
+            
+            if user_input.lower() == 'example':
+                print("\nExample questions:")
+                print("  1. A ball is thrown upward with velocity 20 m/s. What is its maximum height?")
+                print("  2. Calculate the force needed to accelerate a 5kg object at 3 m/s²")
+                print("  3. What is the wavelength of light with frequency 5×10¹⁴ Hz?")
+                print("  4. A 2kg block slides down a 30° incline. What is its acceleration?")
+                continue
+            
+            if user_input.lower() == 'config':
+                print(f"\nCurrent settings:")
+                print(f"  max_gen_len: {max_gen_len}")
+                print(f"  temperature: {temperature}")
+                print(f"  top_p: {top_p}")
+                
+                try:
+                    new_max_len = input(f"\nNew max_gen_len [{max_gen_len}]: ").strip()
+                    if new_max_len:
+                        max_gen_len = int(new_max_len)
+                    
+                    new_temp = input(f"New temperature [{temperature}]: ").strip()
+                    if new_temp:
+                        temperature = float(new_temp)
+                    
+                    new_top_p = input(f"New top_p [{top_p}]: ").strip()
+                    if new_top_p:
+                        top_p = float(new_top_p)
+                    
+                    print("\n✓ Configuration updated!")
+                except ValueError:
+                    print("\n✗ Invalid input. Configuration unchanged.")
+                continue
+            
+            # Generate response
+            full_text, is_finished = generate_response(
+                model=model,
+                tokenizer=tokenizer,
+                question=user_input,
+                device=device,
+                dtype=dtype,
+                max_gen_len=max_gen_len,
+                temperature=temperature,
+                top_p=top_p,
+            )
+            
+            # Print response
+            print_response(full_text)
+            
+            if not is_finished:
+                print("⚠️  Note: Response was truncated (reached max_gen_len)")
+        
+        except KeyboardInterrupt:
+            print("\n\nInterrupted. Type 'quit' to exit.\n")
+            continue
+        except Exception as e:
+            print(f"\n✗ Error: {e}\n")
+            continue
+
+
+def batch_inference_mode(model, tokenizer, device, dtype, config, questions_file, output_file):
+    """Run batch inference on a file of questions."""
+    print(f"\nRunning batch inference on {questions_file}...")
+    
+    max_gen_len = config["training"].get("max_gen_len", 512)
+    
+    # Read questions
+    with open(questions_file, 'r') as f:
+        questions = [line.strip() for line in f if line.strip()]
+    
+    print(f"Found {len(questions)} questions")
+    
+    results = []
+    for i, question in enumerate(questions, 1):
+        print(f"\n[{i}/{len(questions)}] Processing: {question[:60]}...")
+        
+        full_text, is_finished = generate_response(
+            model=model,
+            tokenizer=tokenizer,
+            question=question,
+            device=device,
+            dtype=dtype,
+            max_gen_len=max_gen_len,
+            temperature=0.7,
+            top_p=0.9,
+        )
+        
+        answer = extract_answer(full_text)
+        
+        results.append({
+            'question': question,
+            'full_response': full_text,
+            'answer': answer,
+            'is_finished': is_finished,
+        })
+    
+    # Save results
+    import json
+    with open(output_file, 'w') as f:
+        json.dump(results, f, indent=2)
+    
+    print(f"\n✓ Results saved to {output_file}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Interactive inference for physics problem solver")
+    parser.add_argument("--config", type=str, required=True, help="Path to config YAML file")
+    parser.add_argument("--checkpoint", type=str, help="Path to model checkpoint (optional)")
+    parser.add_argument("--batch", action="store_true", help="Run batch inference mode")
+    parser.add_argument("--questions", type=str, help="Path to questions file (for batch mode)")
+    parser.add_argument("--output", type=str, default="results.json", help="Output file (for batch mode)")
+    
+    args = parser.parse_args()
+    
+    # Load model and tokenizer
+    print("Loading model and tokenizer...")
+    model, tokenizer, device, dtype, config = load_model_and_tokenizer(
+        args.config, 
+        args.checkpoint
+    )
+    print("✓ Model loaded successfully!\n")
+    
+    if args.batch:
+        if not args.questions:
+            print("Error: --questions file required for batch mode")
+            return
+        batch_inference_mode(model, tokenizer, device, dtype, config, args.questions, args.output)
+    else:
+        interactive_mode(model, tokenizer, device, dtype, config)
+
+
+if __name__ == "__main__":
+    main()

questions.txt

@@ -0,0 +1,7 @@
+A train moves 500 kilometers in 5 hours. What is its speed in meters per second?
+Name the following organic compound: CH3CH2CH2OH.
+Implement a python code that outputs the shortest path in a graph
+Proof the Pythagoras theorem
+Current Position Details:White Pieces: King on h1, Rook on e1. Black Pieces: King on h8, Rook on g8. The Challenge: It is White's turn to move. Find the move that results in checkmate (Mate in 1).
+Facing a $8 bet into a $20 pot on the $A♠ 8♣ 4♠$ flop holding $K♠ 9♠$, your pot odds are $28:8$ (or $3.5$-to-$1$), which is slightly less than the $4.1$-to-$1$ needed for a pure $9$-out flush draw, but the call is justified by implied odds since you can expect to win a larger bet if you hit your hand.
+