import osimport sysfrom threading import Threadimport torchimport timefrom transformers import ( AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList) # =========================================================================# 1. CONFIGURATION & MODEL LOADING# =========================================================================REPO_ID = "Phase-Technologies/qwen2.5-3b-claude-distilled-reasoning-dpo" print(f"[*] Hardware Status: CUDA Available: {torch.cuda.is_available()}") tokenizer = AutoTokenizer.from_pretrained(REPO_ID)model = AutoModelForCausalLM.from_pretrained( REPO_ID, torch_dtype=torch.float16, device_map="auto", attn_implementation="sdpa") # =========================================================================# 2. ENHANCED INFERENCE ENGINE# =========================================================================def analyze_inference(prompt): messages = [ {"role": "system", "content": "You are a reasoning assistant. Solve the problem step-by-step and provide a final answer in a box."}, {"role": "user", "content": prompt} ] formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device) streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) # Increased repetition penalty to 1.3 to stop 'Implication:' loops # Added stop_strings for common hallucination patterns gen_kwargs = dict( **inputs, streamer=streamer, max_new_tokens=400, do_sample=True, temperature=0.4, top_p=0.9, repetition_penalty=1.3, stop_strings=["Implication:", "<|im_end|>", "###"], tokenizer=tokenizer, pad_token_id=tokenizer.eos_token_id ) print(f"\n--- TESTING IMPROVED PARAMETERS ---") start_time = time.time() thread = Thread(target=model.generate, kwargs=gen_kwargs) thread.start() generated_text = "" for new_text in streamer: print(new_text, end="", flush=True) generated_text += new_text duration = time.time() - start_time print(f"\n\n[Metric] Speed: {len(tokenizer.encode(generated_text))/duration:.2f} tokens/sec") analyze_inference("Sally has 3 brothers. Each of her brothers has 2 sisters. How many sisters does Sally have?")