Create test_model.py

test_model.py

@@ -0,0 +1,285 @@
+#!/usr/bin/env python3
+"""
+Test script for Mineral Nano 1 Vision AI
+Simple chatbot interface to test your model
+"""
+
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor
+from PIL import Image
+import torch
+
+# ============ CONFIGURATION ============
+MODEL_NAME = "Luke-Bergen/mineral-nano-1"  # Your model
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+# =======================================
+
+def load_model():
+    """Load the model and tokenizer"""
+    print("="*60)
+    print("Loading Mineral Nano 1 Vision Model...")
+    print("="*60)
+    
+    try:
+        print(f"\n📦 Loading from: {MODEL_NAME}")
+        print(f"💻 Device: {DEVICE}")
+        
+        # Load tokenizer
+        print("\n[1/3] Loading tokenizer...")
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        
+        # Load processor (for images)
+        print("[2/3] Loading processor...")
+        processor = AutoProcessor.from_pretrained(MODEL_NAME)
+        
+        # Load model
+        print("[3/3] Loading model (this may take a minute)...")
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            torch_dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32,
+            device_map="auto" if DEVICE == "cuda" else None,
+            low_cpu_mem_usage=True
+        )
+        
+        if DEVICE == "cpu":
+            model = model.to(DEVICE)
+        
+        print("\n✅ Model loaded successfully!")
+        print(f"📊 Parameters: ~2.7B")
+        print(f"🧠 Memory usage: ~{torch.cuda.memory_allocated() / 1e9:.2f} GB" if DEVICE == "cuda" else "")
+        
+        return model, tokenizer, processor
+        
+    except Exception as e:
+        print(f"\n❌ Error loading model: {e}")
+        print("\nPossible issues:")
+        print("  - Model still uploading to HuggingFace")
+        print("  - Model is private (make it public in settings)")
+        print("  - Not enough RAM/VRAM")
+        return None, None, None
+
+
+def chat_text_only(model, tokenizer):
+    """Simple text-only chatbot"""
+    print("\n" + "="*60)
+    print("💬 TEXT CHAT MODE")
+    print("="*60)
+    print("Type your message and press Enter")
+    print("Type 'exit' or 'quit' to stop")
+    print("Type 'clear' to reset conversation")
+    print("="*60 + "\n")
+    
+    conversation_history = []
+    
+    while True:
+        # Get user input
+        user_input = input("You: ").strip()
+        
+        if user_input.lower() in ['exit', 'quit']:
+            print("\n👋 Goodbye!")
+            break
+        
+        if user_input.lower() == 'clear':
+            conversation_history = []
+            print("\n🔄 Conversation cleared!\n")
+            continue
+        
+        if not user_input:
+            continue
+        
+        # Add to conversation
+        conversation_history.append(f"User: {user_input}")
+        
+        # Format prompt
+        prompt = "\n".join(conversation_history) + "\nAssistant:"
+        
+        # Tokenize
+        inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+        
+        # Generate response
+        print("Assistant: ", end="", flush=True)
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=150,
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        # Decode response
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract only the new response
+        response = response[len(prompt):].strip()
+        
+        print(response + "\n")
+        
+        # Add to history
+        conversation_history.append(f"Assistant: {response}")
+        
+        # Keep conversation history manageable
+        if len(conversation_history) > 10:
+            conversation_history = conversation_history[-10:]
+
+
+def chat_with_image(model, processor):
+    """Chat about an image"""
+    print("\n" + "="*60)
+    print("🖼️  IMAGE CHAT MODE")
+    print("="*60)
+    
+    # Get image path
+    image_path = input("Enter image path (or 'back' to return): ").strip()
+    
+    if image_path.lower() == 'back':
+        return
+    
+    try:
+        # Load image
+        image = Image.open(image_path)
+        print(f"✅ Image loaded: {image.size}")
+        
+        print("\nAsk questions about the image (type 'back' to change image)")
+        print("="*60 + "\n")
+        
+        while True:
+            question = input("You: ").strip()
+            
+            if question.lower() == 'back':
+                break
+            
+            if not question:
+                continue
+            
+            # Format prompt with image token
+            prompt = f"<image>{question}"
+            
+            # Process inputs
+            inputs = processor(text=prompt, images=image, return_tensors="pt")
+            inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+            
+            # Generate
+            print("Assistant: ", end="", flush=True)
+            
+            with torch.no_grad():
+                outputs = model.generate(
+                    **inputs,
+                    max_new_tokens=200,
+                    temperature=0.7,
+                    top_p=0.9,
+                    do_sample=True
+                )
+            
+            # Decode
+            response = processor.decode(outputs[0], skip_special_tokens=True)
+            print(response + "\n")
+            
+    except FileNotFoundError:
+        print(f"❌ Image not found: {image_path}")
+    except Exception as e:
+        print(f"❌ Error: {e}")
+
+
+def quick_test(model, tokenizer):
+    """Run a quick test to verify model works"""
+    print("\n" + "="*60)
+    print("🧪 QUICK TEST")
+    print("="*60)
+    
+    test_prompts = [
+        "Hello! What can you do?",
+        "What is 2+2?",
+        "Tell me a short joke."
+    ]
+    
+    for prompt in test_prompts:
+        print(f"\n💬 Test: {prompt}")
+        print("🤖 Response: ", end="", flush=True)
+        
+        inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=50,
+                temperature=0.7,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response = response[len(prompt):].strip()
+        print(response)
+    
+    print("\n" + "="*60)
+    print("✅ Quick test completed!")
+    print("="*60)
+
+
+def main():
+    """Main menu"""
+    print("\n")
+    print("╔════════════════════════════════════════════════════════╗")
+    print("║                                                        ║")
+    print("║          MINERAL NANO 1 VISION - TEST SUITE           ║")
+    print("║              Mid-Range AI Model (2.7B)                ║")
+    print("║                                                        ║")
+    print("╚════════════════════════════════════════════════════════╝")
+    
+    # Load model
+    model, tokenizer, processor = load_model()
+    
+    if model is None:
+        print("\n❌ Failed to load model. Exiting.")
+        return
+    
+    # Main menu loop
+    while True:
+        print("\n" + "="*60)
+        print("MAIN MENU")
+        print("="*60)
+        print("1. 💬 Text Chat (no images)")
+        print("2. 🖼️  Image Chat (describe images)")
+        print("3. 🧪 Quick Test (verify model works)")
+        print("4. 📊 Model Info")
+        print("5. 🚪 Exit")
+        print("="*60)
+        
+        choice = input("\nSelect option (1-5): ").strip()
+        
+        if choice == "1":
+            chat_text_only(model, tokenizer)
+        elif choice == "2":
+            if processor is None:
+                print("❌ Processor not loaded. Cannot use image mode.")
+            else:
+                chat_with_image(model, processor)
+        elif choice == "3":
+            quick_test(model, tokenizer)
+        elif choice == "4":
+            print("\n" + "="*60)
+            print("MODEL INFORMATION")
+            print("="*60)
+            print(f"Model: {MODEL_NAME}")
+            print(f"Parameters: ~2.7B")
+            print(f"Architecture: Vision-Language Model")
+            print(f"Context Length: 8192 tokens")
+            print(f"Image Resolution: 448x448")
+            print(f"Device: {DEVICE}")
+            print("="*60)
+        elif choice == "5":
+            print("\n👋 Goodbye!")
+            break
+        else:
+            print("❌ Invalid option. Please choose 1-5.")
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\n\n👋 Interrupted. Goodbye!")
+    except Exception as e:
+        print(f"\n❌ Unexpected error: {e}")