PredictiveManish/Trimurti-LM
019
1"""
2Step 4: Test your trained multilingual model
3"""
4
5import torch
6from transformers import GPT2LMHeadModel
7import sentencepiece as spm
8import os
9from pathlib import Path
10
11class MultilingualModel:
12 def __init__(self, model_path="./checkpoints_tiny/final"):
13 print("="*60)
14 print("LOADING MULTILINGUAL MODEL")
15 print("="*60)
16
17 # Check if model exists
18 if not os.path.exists(model_path):
19 print(f"❌ Model not found at: {model_path}")
20 print("Available checkpoints:")
21 checkpoints = list(Path("./checkpoints_tiny").glob("checkpoint-*"))
22 checkpoints += list(Path("./checkpoints_tiny").glob("step*"))
23 checkpoints += list(Path("./checkpoints_tiny").glob("final"))
24
25 for cp in checkpoints:
26 if cp.is_dir():
27 print(f" - {cp}")
28
29 if checkpoints:
30 model_path = str(checkpoints[-1]) # Use most recent
31 print(f"Using: {model_path}")
32 else:
33 raise FileNotFoundError("No checkpoints found!")
34
35 # Load tokenizer
36 tokenizer_path = os.path.join(model_path, "tokenizer", "spiece.model")
37 if not os.path.exists(tokenizer_path):
38 tokenizer_path = "./final_corpus/multilingual_spm.model"
39
40 print(f"Loading tokenizer from: {tokenizer_path}")
41 self.tokenizer = spm.SentencePieceProcessor()
42 self.tokenizer.load(tokenizer_path)
43
44 # Load model
45 print(f"Loading model from: {model_path}")
46 self.model = GPT2LMHeadModel.from_pretrained(model_path)
47
48 # Setup device
49 self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
50 self.model.to(self.device)
51 self.model.eval()
52
53 print(f"✅ Model loaded on: {self.device}")
54 print(f" Parameters: {sum(p.numel() for p in self.model.parameters())/1e6:.1f}M")
55 print("="*60)
56
57 def generate(self, prompt, max_length=100, temperature=0.7, top_k=50, top_p=0.95):
58 """Generate text from prompt"""
59 # Add language tag if missing
60 if not any(prompt.startswith(tag) for tag in ['[EN]', '[HI]', '[PA]']):
61 # Try to detect language
62 if any(char in prompt for char in 'अआइईउऊएऐओऔकखगघचछजझटठडढणतथदधनपफबभमयरलवशषसह'):
63 prompt = f"[HI] {prompt}"
64 elif any(char in prompt for char in 'ਅਆਇਈਉਊਏਐਓਔਕਖਗਘਚਛਜਝਟਠਡਢਣਤਥਦਧਨਪਫਬਭਮਯਰਲਵਸ਼ਸਹ'):
65 prompt = f"[PA] {prompt}"
66 else:
67 prompt = f"[EN] {prompt}"
68
69 # Encode
70 input_ids = self.tokenizer.encode(prompt)
71 input_tensor = torch.tensor([input_ids], device=self.device)
72
73 # Generate
74 with torch.no_grad():
75 output = self.model.generate(
76 input_ids=input_tensor,
77 max_length=max_length,
78 temperature=temperature,
79 do_sample=True,
80 top_k=top_k,
81 top_p=top_p,
82 pad_token_id=self.tokenizer.pad_id() if self.tokenizer.pad_id() > 0 else 0,
83 eos_token_id=self.tokenizer.eos_id() if self.tokenizer.eos_id() > 0 else 2,
84 repetition_penalty=1.1,
85 )
86
87 # Decode
88 generated = self.tokenizer.decode(output[0].tolist())
89
90 # Clean up (remove prompt if it's repeated)
91 if generated.startswith(prompt):
92 result = generated[len(prompt):].strip()
93 else:
94 result = generated
95
96 return result
97
98 def batch_generate(self, prompts, **kwargs):
99 """Generate for multiple prompts"""
100 results = []
101 for prompt in prompts:
102 result = self.generate(prompt, **kwargs)
103 results.append(result)
104 return results
105
106 def calculate_perplexity(self, text):
107 """Calculate perplexity of given text"""
108 input_ids = self.tokenizer.encode(text)
109 if len(input_ids) < 2:
110 return float('inf')
111
112 input_tensor = torch.tensor([input_ids], device=self.device)
113
114 with torch.no_grad():
115 outputs = self.model(input_ids=input_tensor, labels=input_tensor)
116 loss = outputs.loss
117
118 perplexity = torch.exp(loss).item()
119 return perplexity
120
121 def interactive_mode(self):
122 """Interactive chat with model"""
123 print("\n" + "="*60)
124 print("INTERACTIVE MODE")
125 print("="*60)
126 print("Enter prompts in any language (add [EN], [HI], [PA] tags)")
127 print("Commands: /temp X, /len X, /quit, /help")
128 print("="*60)
129
130 temperature = 0.7
131 max_length = 100
132
133 while True:
134 try:
135 user_input = input("\nYou: ").strip()
136
137 if not user_input:
138 continue
139
140 # Handle commands
141 if user_input.startswith('/'):
142 if user_input == '/quit':
143 break
144 elif user_input == '/help':
145 print("Commands:")
146 print(" /temp X - Set temperature (0.1 to 2.0)")
147 print(" /len X - Set max length (20 to 500)")
148 print(" /quit - Exit")
149 print(" /help - Show this help")
150 continue
151 elif user_input.startswith('/temp'):
152 try:
153 temp = float(user_input.split()[1])
154 if 0.1 <= temp <= 2.0:
155 temperature = temp
156 print(f"Temperature set to {temperature}")
157 else:
158 print("Temperature must be between 0.1 and 2.0")
159 except:
160 print("Usage: /temp 0.7")
161 continue
162 elif user_input.startswith('/len'):
163 try:
164 length = int(user_input.split()[1])
165 if 20 <= length <= 500:
166 max_length = length
167 print(f"Max length set to {max_length}")
168 else:
169 print("Length must be between 20 and 500")
170 except:
171 print("Usage: /len 100")
172 continue
173
174 # Generate response
175 print("Model: ", end="", flush=True)
176 response = self.generate(user_input, max_length=max_length, temperature=temperature)
177 print(response)
178
179 except KeyboardInterrupt:
180 print("\n\nExiting...")
181 break
182 except Exception as e:
183 print(f"Error: {e}")
184
185def run_tests():
186 """Run comprehensive tests"""
187 print("\n" + "="*60)
188 print("COMPREHENSIVE MODEL TESTS")
189 print("="*60)
190
191 # Load model
192 model = MultilingualModel()
193
194 # Test prompts by language
195 test_suites = {
196 "English": [
197 "[EN] The weather today is",
198 "[EN] I want to learn",
199 "[EN] Artificial intelligence",
200 "[EN] The capital of India is",
201 "[EN] Once upon a time",
202 ],
203 "Hindi": [
204 "[HI] आज का मौसम",
205 "[HI] मैं सीखना चाहता हूं",
206 "[HI] कृत्रिम बुद्धिमत्ता",
207 "[HI] भारत की राजधानी है",
208 "[HI] एक बार की बात है",
209 ],
210 "Punjabi": [
211 "[PA] ਅੱਜ ਦਾ ਮੌਸਮ",
212 "[PA] ਮੈਂ ਸਿੱਖਣਾ ਚਾਹੁੰਦਾ ਹਾਂ",
213 "[PA] ਕ੍ਰਿਤਰਿਮ ਬੁੱਧੀ",
214 "[PA] ਭਾਰਤ ਦੀ ਰਾਜਧਾਨੀ ਹੈ",
215 "[PA] ਇੱਕ ਵਾਰ ਦੀ ਗੱਲ ਹੈ",
216 ],
217 "Language Switching": [
218 "[EN] Hello [HI] नमस्ते",
219 "[HI] यह अच्छा है [EN] this is good",
220 "[PA] ਸਤਿ ਸ੍ਰੀ ਅਕਾਲ [EN] Hello everyone",
221 ],
222 "Code Mixing": [
223 "Hello दुनिया", # No tag, should auto-detect
224 "मेरा name है", # Hindi + English
225 "Today मौसम is good", # English + Hindi
226 ]
227 }
228
229 for suite_name, prompts in test_suites.items():
230 print(f"\n{'='*40}")
231 print(f"{suite_name.upper()} TESTS")
232 print('='*40)
233
234 for i, prompt in enumerate(prompts):
235 print(f"\nTest {i+1}:")
236 print(f"Prompt: {prompt}")
237
238 # Generate
239 response = model.generate(prompt, max_length=50, temperature=0.7)
240 print(f"Response: {response}")
241
242 # Calculate perplexity
243 try:
244 perplexity = model.calculate_perplexity(response)
245 print(f"Perplexity: {perplexity:.2f}")
246 except:
247 pass
248
249 print("-" * 40)
250
251def benchmark_model():
252 """Benchmark model performance"""
253 print("\n" + "="*60)
254 print("MODEL BENCHMARK")
255 print("="*60)
256
257 model = MultilingualModel()
258
259 import time
260
261 # Test generation speed
262 test_prompt = "[EN] The quick brown fox"
263
264 times = []
265 for _ in range(10):
266 start = time.time()
267 model.generate(test_prompt, max_length=50)
268 end = time.time()
269 times.append(end - start)
270
271 avg_time = sum(times) / len(times)
272 print(f"Average generation time (50 tokens): {avg_time:.3f}s")
273 print(f"Tokens per second: {50/avg_time:.1f}")
274
275 # Memory usage
276 if torch.cuda.is_available():
277 memory_allocated = torch.cuda.memory_allocated() / 1e9
278 memory_reserved = torch.cuda.memory_reserved() / 1e9
279 print(f"GPU Memory allocated: {memory_allocated:.2f} GB")
280 print(f"GPU Memory reserved: {memory_reserved:.2f} GB")
281
282def create_web_interface():
283 """Simple web interface for the model"""
284 html_code = """
285<!DOCTYPE html>
286<html>
287<head>
288 <title>Multilingual LM Demo</title>
289 <style>
290 body { font-family: Arial, sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; }
291 .container { display: flex; flex-direction: column; gap: 20px; }
292 textarea { width: 100%; height: 100px; padding: 10px; font-size: 16px; }
293 button { padding: 10px 20px; background: #4CAF50; color: white; border: none; cursor: pointer; }
294 button:hover { background: #45a049; }
295 .output { border: 1px solid #ccc; padding: 15px; min-height: 100px; background: #f9f9f9; }
296 .language-tag { display: inline-block; margin: 5px; padding: 5px 10px; background: #e0e0e0; cursor: pointer; }
297 </style>
298</head>
299<body>
300 <div class="container">
301 <h1>Multilingual Language Model Demo</h1>
302
303 <div>
304 <strong>Language:</strong>
305 <span class="language-tag" onclick="setLanguage('[EN] ')">English</span>
306 <span class="language-tag" onclick="setLanguage('[HI] ')">Hindi</span>
307 <span class="language-tag" onclick="setLanguage('[PA] ')">Punjabi</span>
308 </div>
309
310 <textarea id="prompt" placeholder="Enter your prompt here..."></textarea>
311
312 <div>
313 <label>Temperature: <input type="range" id="temp" min="0.1" max="2.0" step="0.1" value="0.7"></label>
314 <label>Max Length: <input type="number" id="maxlen" min="20" max="500" value="100"></label>
315 </div>
316
317 <button onclick="generate()">Generate</button>
318
319 <div class="output" id="output">Response will appear here...</div>
320 </div>
321
322 <script>
323 function setLanguage(tag) {
324 document.getElementById('prompt').value = tag;
325 }
326
327 async function generate() {
328 const prompt = document.getElementById('prompt').value;
329 const temp = document.getElementById('temp').value;
330 const maxlen = document.getElementById('maxlen').value;
331
332 document.getElementById('output').innerHTML = 'Generating...';
333
334 try {
335 const response = await fetch('/generate', {
336 method: 'POST',
337 headers: {'Content-Type': 'application/json'},
338 body: JSON.stringify({prompt, temp, maxlen})
339 });
340
341 const data = await response.json();
342 document.getElementById('output').innerHTML = data.response;
343 } catch (error) {
344 document.getElementById('output').innerHTML = 'Error: ' + error;
345 }
346 }
347 </script>
348</body>
349</html>
350 """
351
352 # Save HTML
353 with open("model_demo.html", "w", encoding="utf-8") as f:
354 f.write(html_code)
355
356 print("Web interface saved as model_demo.html")
357 print("To use it, you need a backend server (see create_server.py)")
358
359def main():
360 """Main function"""
361 print("\n" + "="*60)
362 print("MULTILINGUAL MODEL PLAYGROUND")
363 print("="*60)
364 print("\nOptions:")
365 print("1. Interactive chat")
366 print("2. Run comprehensive tests")
367 print("3. Benchmark model")
368 print("4. Create web interface")
369 print("5. Quick generation test")
370 print("6. Exit")
371
372 # Load model once
373 model = None
374
375 while True:
376 try:
377 choice = input("\nSelect option (1-6): ").strip()
378
379 if choice == '1':
380 if model is None:
381 model = MultilingualModel()
382 model.interactive_mode()
383
384 elif choice == '2':
385 run_tests()
386
387 elif choice == '3':
388 benchmark_model()
389
390 elif choice == '4':
391 create_web_interface()
392
393 elif choice == '5':
394 if model is None:
395 model = MultilingualModel()
396
397 prompt = input("Enter prompt: ").strip()
398 if prompt:
399 response = model.generate(prompt)
400 print(f"\nResponse: {response}")
401
402 elif choice == '6':
403 print("Goodbye!")
404 break
405
406 else:
407 print("Invalid choice. Please enter 1-6.")
408
409 except KeyboardInterrupt:
410 print("\n\nExiting...")
411 break
412 except Exception as e:
413 print(f"Error: {e}")
414 import traceback
415 traceback.print_exc()
416
417if __name__ == "__main__":
418 main()