Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -1,23 +1,19 @@
 """
-Test Florence-2-Large model integration with vGPU.
-Configure PyTorch to use vGPU as device and run image inference.
+Test Llama-2-7b-instruct model integration with vGPU.
+Configure PyTorch to use vGPU as device for text generation.
 """
 import logging
 import os
 import time
 from contextlib import contextmanager
-from io import BytesIO
 from typing import Any, Optional
 
 import torch
-from torch import nn
 import torch.nn.functional as F
-from torch.overrides import TorchFunctionMode
-from PIL import Image
 from transformers import (
-    AutoProcessor,
-    AutoModel,
-    AutoConfig
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TextStreamer
 )
 from virtual_vram import VirtualVRAM
 from http_storage import HTTPGPUStorage
@@ -68,33 +64,20 @@ def get_model_size(model):
         buffer_size += buffer.nelement() * buffer.element_size()
     return param_size + buffer_size
 
-def load_image(image_name):
-    """Load and preprocess image from sample_task folder"""
-    try:
-        image_path = os.path.join("sample_task", image_name)
-        if not os.path.exists(image_path):
-            raise FileNotFoundError(f"Image not found: {image_path}")
-            
-        image = Image.open(image_path)
-        # Convert to RGB if needed
-        if image.mode != 'RGB':
-            image = image.convert('RGB')
-        logger.info(f"Loaded image from {image_path}: size={image.size}")
-        return image
-    except Exception as e:
-        logger.error(f"Image loading failed: {str(e)}")
-        raise
+def prepare_prompt(instruction: str) -> str:
+    """Prepare a prompt for Llama-2 using its chat format."""
+    # Format: <s>[INST] instruction [/INST] assistant response </s>[INST] ...
+    return f"<s>[INST] {instruction} [/INST]"
 
 def test_ai_integration_http():
-    """Test Florence-2-Large model on vGPU with image inference"""
-    logger.info("Starting vGPU image inference test")
+    """Test Llama-2-7b-instruct model on vGPU with text generation"""
+    logger.info("Starting vGPU text generation test")
     
     status = {
         'model_loaded': False,
-        'processor_loaded': False,
+        'tokenizer_loaded': False,
         'model_on_vgpu': False,
-        'image_processed': False,
-        'inference_complete': False,
+        'generation_complete': False,
         'cleanup_success': False
     }
     
@@ -108,40 +91,37 @@ def test_ai_integration_http():
             device = setup_vgpu()
             logger.info(f"vGPU initialized with device {device}")
             
-            # Load Florence model and processor
-            model_name = "microsoft/florence-2-large"
+            # Load Llama model and tokenizer
+            model_name = "meta-llama/Llama-2-7b-chat-hf"
             logger.info(f"Loading {model_name}")
             
             try:
                 # Disable transformers logging temporarily
-                import logging
                 transformers_logger = logging.getLogger("transformers")
                 original_level = transformers_logger.level
                 transformers_logger.setLevel(logging.ERROR)
                 
                 try:
-                    # Load processor first
-                    processor = AutoProcessor.from_pretrained(
+                    # Load tokenizer first
+                    tokenizer = AutoTokenizer.from_pretrained(
                         model_name,
-                        trust_remote_code=True
+                        trust_remote_code=True,
+                        use_fast=True
                     )
-                    status['processor_loaded'] = True
+                    status['tokenizer_loaded'] = True
                     
-                    # Import the specific model class
-                    from transformers.models.florence.modeling_florence import Florence2Model
-                    
-                    # Load model directly with specific class
-                    model = Florence2Model.from_pretrained(
+                    # Load model with full precision
+                    model = AutoModelForCausalLM.from_pretrained(
                         model_name,
                         trust_remote_code=True,
-                        torch_dtype=torch.float32,
-                        device_map=None,
-                        ignore_mismatched_sizes=True
+                        torch_dtype=torch.float32,  # Use full precision
+                        device_map=None,  # Don't auto-map devices
+                        use_safetensors=True
                     )
                     status['model_loaded'] = True
                     
                     # Log model details
-                    logger.info(f"Processor type: {type(processor).__name__}")
+                    logger.info(f"Tokenizer type: {type(tokenizer).__name__}")
                     logger.info(f"Model type: {type(model).__name__}")
                     
                     # Log model architecture
@@ -175,62 +155,51 @@ def test_ai_integration_http():
                 logger.error(f"Model transfer to vGPU failed: {str(e)}")
                 raise
             
-            # Prepare image input from sample_task folder
-            try:
-                # Load image from sample_task directory
-                image_name = "sample1.jpg"  # Replace with your image name
-                image = load_image(image_name)
-                
-                # Process image with Florence processor
-                inputs = processor(images=image, return_tensors="pt")
-                if not inputs or 'pixel_values' not in inputs:
-                    raise ValueError("Invalid processor output")
-                
-                # Move inputs to vGPU
-                inputs = {k: to_vgpu(v, vram=vram) for k, v in inputs.items()}
-                status['image_processed'] = True
-                logger.info(f"Image processed: shape={inputs['pixel_values'].shape}")
-            except Exception as e:
-                logger.error(f"Image preparation failed: {str(e)}")
-                raise
-            
-            # Run image inference with monitoring
-            logger.info("Running image inference...")
+            # Run text generation
+            logger.info("Running text generation...")
             start = time.time()
             peak_mem = initial_mem
             
             try:
+                # Prepare input prompt
+                instruction = "Explain how virtual GPUs work in simple terms."
+                prompt = prepare_prompt(instruction)
+                
+                # Tokenize input
+                inputs = tokenizer(prompt, return_tensors="pt")
+                inputs = {k: to_vgpu(v, vram=vram) for k, v in inputs.items()}
+                
+                # Set up streamer for token-by-token output
+                streamer = TextStreamer(tokenizer)
+                
                 with torch.no_grad():
-                    # Get image embeddings
-                    outputs = model(**inputs)
-                    image_features = outputs.last_hidden_state[:, 0]  # Take [CLS] token features
-                    
-                    # Normalize features
-                    image_features = F.normalize(image_features, dim=-1)
+                    # Generate text
+                    outputs = model.generate(
+                        **inputs,
+                        max_length=512,
+                        temperature=0.7,
+                        top_p=0.95,
+                        top_k=40,
+                        num_beams=1,
+                        streamer=streamer,
+                        pad_token_id=tokenizer.pad_token_id
+                    )
                     
                     if hasattr(storage, 'get_used_memory'):
                         peak_mem = max(peak_mem, storage.get_used_memory())
                     
                     inference_time = time.time() - start
-                    status['inference_complete'] = True
+                    status['generation_complete'] = True
                     
                     # Log performance metrics
-                    logger.info(f"Inference stats:")
+                    logger.info(f"\nGeneration stats:")
                     logger.info(f"- Time: {inference_time:.4f}s")
                     logger.info(f"- Memory peak: {(peak_mem - initial_mem)/1e9:.2f} GB")
-                    logger.info(f"- Image features shape: {image_features.shape}")
-                    logger.info(f"- Feature norm: {torch.norm(image_features).item():.4f}")
-                    logger.info(f"- Output device: {image_features.device}")
-                    
-                    # Optionally compute confidence scores
-                    if hasattr(outputs, 'logits'):
-                        logits = outputs.logits
-                        probs = F.softmax(logits, dim=-1)
-                        confidence = torch.max(probs).item()
-                        logger.info(f"- Confidence: {confidence:.4f}")
+                    logger.info(f"- Output length: {len(outputs[0])}")
+                    logger.info(f"- Output device: {outputs.device}")
                         
             except Exception as e:
-                logger.error(f"Image inference failed: {str(e)}")
+                logger.error(f"Text generation failed: {str(e)}")
                 raise
             
         except Exception as e: