Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3from peft import PeftModel
4import warnings
5import json
6warnings.filterwarnings("ignore")
7
8# Load base model and tokenizer (using actual base model from training)
9base_model_name = "unsloth/Llama-3.2-1B-Instruct" # Actual base model used in training
10model_name = "shiprocket-ai/multitask-address-reasoning-llama-1B-model"
11
12print("📥 Loading tokenizer...")
13# Load tokenizer
14tokenizer = AutoTokenizer.from_pretrained(model_name)
15
16# Add pad token if missing
17if tokenizer.pad_token is None:
18 tokenizer.pad_token = tokenizer.eos_token
19
20print("📥 Loading base model...")
21# Load base model (non-quantized version as per training script)
22base_model = AutoModelForCausalLM.from_pretrained(
23 base_model_name,
24 torch_dtype=torch.float16,
25 device_map="auto",
26 trust_remote_code=True
27)
28
29print("📥 Loading LoRA adapter...")
30# Load LoRA adapter
31model = PeftModel.from_pretrained(base_model, model_name)
32
33print("✅ Model loaded successfully!")
34
35def process_address_with_reasoning(prompt, max_new_tokens=400):
36 """Process address with Chain of Thought reasoning (as trained)"""
37
38 # Tokenize
39 inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
40
41 # Move inputs to model device
42 device = next(model.parameters()).device
43 inputs = {k: v.to(device) for k, v in inputs.items()}
44
45 # Generate with reasoning (matching training parameters)
46 with torch.no_grad():
47 outputs = model.generate(
48 **inputs,
49 max_new_tokens=max_new_tokens,
50 temperature=0.1, # Lower temperature as used in training testing
51 do_sample=True,
52 pad_token_id=tokenizer.eos_token_id,
53 use_cache=True
54 )
55
56 # Decode only the new tokens
57 input_length = inputs['input_ids'].shape[1]
58 generated_tokens = outputs[0][input_length:]
59 response = tokenizer.decode(generated_tokens, skip_special_tokens=True)
60
61 return response.strip()
62
63def fix_address_with_reasoning(address, max_new_tokens=400):
64 """Fix address with detailed Chain of Thought reasoning"""
65
66 messages = [
67 {"role": "user", "content": f"Fix and extract components from: {address}"}
68 ]
69
70 prompt = tokenizer.apply_chat_template(
71 messages,
72 tokenize=False,
73 add_generation_prompt=True
74 )
75
76 return process_address_with_reasoning(prompt, max_new_tokens)
77
78def answer_geographic_question(question, max_new_tokens=150):
79 """Answer geographic questions about addresses"""
80
81 messages = [
82 {"role": "user", "content": question}
83 ]
84
85 prompt = tokenizer.apply_chat_template(
86 messages,
87 tokenize=False,
88 add_generation_prompt=True
89 )
90
91 return process_address_with_reasoning(prompt, max_new_tokens)
92
93def extract_components(address, max_new_tokens=200):
94 """Extract address components with reasoning"""
95
96 messages = [
97 {"role": "user", "content": f"Extract all components from this address: {address}"}
98 ]
99
100 prompt = tokenizer.apply_chat_template(
101 messages,
102 tokenize=False,
103 add_generation_prompt=True
104 )
105
106 return process_address_with_reasoning(prompt, max_new_tokens)
107
108# Test cases based on training script examples
109print("""
110🏠 MULTI-TASK ADDRESS MODEL EXAMPLES""")
111print("=" * 60)
112print("""🧠 Testing Chain of Thought reasoning + Geographic Q&A""")
113print("📊 Model trained with LoRA r=64, alpha=128 for complex reasoning")
114print("=" * 60)
115
116# Test address correction with reasoning (exact example from training)
117test_addresses = [
118 "pandit nagla badi masjid moradabad 244001",
119 "sec 14 gurgoan haryana 122001",
120 "koramangala bangalor 560095",
121 "dlf cyber city gurgaon haryana"
122]
123
124print(f"""
125🔧 TESTING ADDRESS CORRECTION WITH CHAIN OF THOUGHT:""")
126print("-" * 50)
127
128for i, test_address in enumerate(test_addresses, 1):
129 print(f"""
130📍 Test {i}: {test_address}""")
131 result = fix_address_with_reasoning(test_address)
132 print(f"🤖 Chain of Thought Response:")
133 print(f" {result}")
134 print("-" * 40)
135
136# Test geographic Q&A (examples from training script)
137qa_tests = [
138 "Which state is Mumbai in?",
139 "What is the pincode of Bangalore?",
140 "Is Delhi a metro city?",
141 "What tier city is Pune?",
142 "Where is Connaught Place located?",
143 "What state does Hyderabad belong to?",
144 "Name a city in Karnataka.",
145 "What is the postal code for Gurgaon?",
146 "Which state is New Delhi in?", # Training example
147 "What cities are in Maharashtra?"
148]
149
150print(f"""
151❓ TESTING GEOGRAPHIC Q&A:""")
152print("-" * 50)
153
154for i, question in enumerate(qa_tests[:8], 1): # Test first 8 questions
155 print(f"""
156❓ Q{i}: {question}""")
157 result = answer_geographic_question(question)
158 print(f"🤖 Answer: {result}")
159
160# Test component extraction
161print(f"""
162📊 TESTING COMPONENT EXTRACTION:""")
163print("-" * 50)
164
165extraction_tests = [
166 "Flat 203, Emerald Heights, Sector 15, Gurugram, Haryana 122001",
167 "DLF Cyber City, Cyber Hub, Gurgaon, Haryana",
168 "Connaught Place, New Delhi, Delhi 110001"
169]
170
171for i, test_address in enumerate(extraction_tests, 1):
172 print(f"""
173📊 Extract {i}: {test_address}""")
174 result = extract_components(test_address)
175 print(f"🤖 Components: {result}")
176
177print(f"""
178✅ ALL TESTS COMPLETED!""")
179print(f"""🧠 Model demonstrates Chain of Thought reasoning""")
180print(f"""📍 Geographic knowledge from NER training data""")
181print(f"""🔧 Address correction with detailed analysis""")Final Training Loss: 0.5506
Training Runtime: 3701.74 seconds (~1 hour)
Training Samples/Second: 3.749
Training Steps/Second: 0.118
Total Epochs: 3.0<|begin_of_text|><|start_header_id|>user<|end_header_id|>
Fix and extract components from: [address]<|eot_id|><|start_header_id|>assistant<|end_header_id|>
<|begin_of_text|><|start_header_id|>user<|end_header_id|>
Which state is [location] in?<|eot_id|><|start_header_id|>assistant<|end_header_id|>
<|begin_of_text|><|start_header_id|>user<|end_header_id|>
Extract all components from this address: [address]<|eot_id|><|start_header_id|>assistant<|end_header_id|>
adapter_config.json: LoRA adapter configurationadapter_model.safetensors: LoRA adapter weightstokenizer_config.json: Tokenizer configurationtokenizer.json: Tokenizer vocabulary and settingsspecial_tokens_map.json: Special tokens mappingchat_template.jinja: Chat template for conversations1@misc{multitask-address-reasoning-model,
2 title={Multi-Task Address Reasoning Model},
3 year={2025},
4 publisher={Hugging Face},
5 url={https://huggingface.co/shiprocket-ai/multitask-address-reasoning-llama-1B-model}
6}