Update app.py

app.py

@@ -3,7 +3,6 @@ import torch
 import sys
 import traceback
 import os
-from huggingface_hub import hf_hub_download
 
 def system_info():
     try:
@@ -27,68 +26,171 @@ def system_info():
     except Exception as e:
         return f"Error: {str(e)}\n\n{traceback.format_exc()}"
 
-def test_gemma_gguf():
+def test_phi3_mini():
     try:
         result = []
-        result.append("Testing Gemma 3 GGUF model...")
-        
-        # First, check if llama-cpp-python is installed
-        try:
-            import llama_cpp
-            result.append(f"llama_cpp version: {llama_cpp.__version__}")
-        except ImportError:
-            result.append("llama-cpp-python not installed. Installing now...")
-            import subprocess
-            subprocess.check_call([sys.executable, "-m", "pip", "install", "llama-cpp-python"])
-            import llama_cpp
-            result.append(f"llama_cpp version: {llama_cpp.__version__}")
-        
-        # Download the model if not already downloaded
-        model_id = "google/gemma-3-27b-it-qat-q4_0-gguf"
-        model_filename = "gemma-3-27b-it-qat-q4_0.gguf"
-        
-        result.append(f"Downloading {model_id} if not already present...")
-        model_path = hf_hub_download(
-            repo_id=model_id,
-            filename=model_filename,
-            resume_download=True
-        )
-        result.append(f"Model downloaded to: {model_path}")
+        result.append("Testing Phi-3 Mini model...")
+        
+        # Use Phi-3 Mini model with 4-bit quantization
+        model_id = "microsoft/Phi-3-mini-4k-instruct"
+        
+        result.append(f"Loading tokenizer from {model_id}...")
+        from transformers import AutoTokenizer, AutoModelForCausalLM
         
-        # Load the model
-        result.append("Loading model...")
-        from llama_cpp import Llama
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
         
-        llm = Llama(
-            model_path=model_path,
-            n_ctx=2048,  # Context window size
-            n_gpu_layers=-1  # Use all available GPU layers
+        result.append("Loading model with quantization...")
+        from transformers import BitsAndBytesConfig
+        
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_quant_type="nf4"
+        )
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            quantization_config=quantization_config,
+            device_map="auto"
         )
         
-        # Generate text
         result.append("Generating text...")
         prompt = "Write a short poem about artificial intelligence."
         
-        output = llm(
-            prompt,
-            max_tokens=100,
-            temperature=0.7,
-            top_p=0.95,
-            echo=False
+        # Format prompt for Phi-3
+        messages = [
+            {"role": "user", "content": prompt}
+        ]
+        prompt = tokenizer.apply_chat_template(messages, tokenize=False)
+        
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        outputs = model.generate(**inputs, max_new_tokens=100)
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract the assistant's response
+        if "<assistant>" in generated_text and "</assistant>" in generated_text:
+            response = generated_text.split("<assistant>")[1].split("</assistant>")[0].strip()
+        else:
+            response = generated_text.replace(prompt, "").strip()
+        
+        result.append(f"Generated text: {response}")
+        result.append("Phi-3 Mini test successful!")
+        
+        return "\n".join(result)
+    except Exception as e:
+        return f"Error: {str(e)}\n\n{traceback.format_exc()}"
+
+def test_image_classification():
+    try:
+        result = []
+        result.append("Testing image classification...")
+        
+        # Use a lightweight vision model
+        from transformers import AutoImageProcessor, AutoModelForImageClassification
+        
+        result.append("Loading image processor and model...")
+        processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50")
+        model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-50")
+        
+        result.append("Loading test image...")
+        import requests
+        from PIL import Image
+        from io import BytesIO
+        
+        response = requests.get("https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg")
+        img = Image.open(BytesIO(response.content))
+        
+        result.append("Processing image...")
+        inputs = processor(images=img, return_tensors="pt")
+        outputs = model(**inputs)
+        
+        # Get predicted class
+        predicted_class_idx = outputs.logits.argmax(-1).item()
+        predicted_class = model.config.id2label[predicted_class_idx]
+        
+        result.append(f"Predicted class: {predicted_class}")
+        result.append("Image classification test successful!")
+        
+        return "\n".join(result)
+    except Exception as e:
+        return f"Error: {str(e)}\n\n{traceback.format_exc()}"
+
+def test_phi3_with_image():
+    try:
+        result = []
+        result.append("Testing Phi-3 Mini with image description...")
+        
+        # First, classify the image
+        from transformers import AutoImageProcessor, AutoModelForImageClassification
+        import requests
+        from PIL import Image
+        from io import BytesIO
+        
+        result.append("Loading image and classifying it...")
+        img_processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50")
+        img_model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-50")
+        
+        response = requests.get("https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg")
+        img = Image.open(BytesIO(response.content))
+        
+        inputs = img_processor(images=img, return_tensors="pt")
+        outputs = img_model(**inputs)
+        
+        predicted_class_idx = outputs.logits.argmax(-1).item()
+        predicted_class = img_model.config.id2label[predicted_class_idx]
+        
+        result.append(f"Image classified as: {predicted_class}")
+        
+        # Now use Phi-3 to describe the image based on the classification
+        result.append("Loading Phi-3 Mini model...")
+        from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+        
+        model_id = "microsoft/Phi-3-mini-4k-instruct"
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_quant_type="nf4"
+        )
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            quantization_config=quantization_config,
+            device_map="auto"
         )
         
-        generated_text = output["choices"][0]["text"]
-        result.append(f"Generated text: {generated_text}")
-        result.append("Gemma 3 GGUF test successful!")
+        # Create a prompt that includes the image classification
+        prompt = f"I have an image that contains a {predicted_class}. Please write a detailed description of what this might look like, and explain some interesting facts about {predicted_class}."
+        
+        # Format prompt for Phi-3
+        messages = [
+            {"role": "user", "content": prompt}
+        ]
+        formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False)
+        
+        result.append("Generating description based on image classification...")
+        inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
+        outputs = model.generate(**inputs, max_new_tokens=200)
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract the assistant's response
+        if "<assistant>" in generated_text and "</assistant>" in generated_text:
+            response = generated_text.split("<assistant>")[1].split("</assistant>")[0].strip()
+        else:
+            response = generated_text.replace(formatted_prompt, "").strip()
+        
+        result.append(f"Generated description: {response}")
+        result.append("Phi-3 with image test successful!")
         
         return "\n".join(result)
     except Exception as e:
         return f"Error: {str(e)}\n\n{traceback.format_exc()}"
 
 # Create Gradio interface
-with gr.Blocks(title="Gemma 3 GGUF Test") as demo:
-    gr.Markdown("# Gemma 3 GGUF Test")
-    gr.Markdown("Testing Gemma 3 model in GGUF format using llama-cpp-python.")
+with gr.Blocks(title="StaffManager AI Assistant") as demo:
+    gr.Markdown("# StaffManager AI Assistant")
+    gr.Markdown("Testing open-access models for text and image processing.")
     
     with gr.Tab("System Info"):
         with gr.Row():
@@ -103,18 +205,57 @@ with gr.Blocks(title="Gemma 3 GGUF Test") as demo:
             outputs=[info_result]
         )
     
-    with gr.Tab("Gemma 3 GGUF Test"):
+    with gr.Tab("Text Generation"):
+        with gr.Row():
+            with gr.Column():
+                phi3_button = gr.Button("Generate Text with Phi-3 Mini")
+            with gr.Column():
+                phi3_result = gr.Textbox(label="Generated Text", lines=20)
+        
+        phi3_button.click(
+            fn=test_phi3_mini,
+            inputs=[],
+            outputs=[phi3_result]
+        )
+    
+    with gr.Tab("Image Classification"):
+        with gr.Row():
+            with gr.Column():
+                image_button = gr.Button("Classify Sample Image")
+            with gr.Column():
+                image_result = gr.Textbox(label="Classification Results", lines=20)
+        
+        image_button.click(
+            fn=test_image_classification,
+            inputs=[],
+            outputs=[image_result]
+        )
+    
+    with gr.Tab("Image Description"):
         with gr.Row():
             with gr.Column():
-                gemma_button = gr.Button("Test Gemma 3 GGUF")
+                combined_button = gr.Button("Generate Image Description")
             with gr.Column():
-                gemma_result = gr.Textbox(label="Test Results", lines=20)
+                combined_result = gr.Textbox(label="Description Results", lines=20)
         
-        gemma_button.click(
-            fn=test_gemma_gguf,
+        combined_button.click(
+            fn=test_phi3_with_image,
             inputs=[],
-            outputs=[gemma_result]
+            outputs=[combined_result]
         )
+    
+    with gr.Tab("About"):
+        gr.Markdown("""
+        ## About StaffManager AI Assistant
+        
+        This Space demonstrates AI capabilities for StaffManager using open-access models:
+        
+        - **Text Generation**: Uses Microsoft's Phi-3 Mini model
+        - **Image Classification**: Uses Microsoft's ResNet-50 model
+        - **Image Description**: Combines both models to classify and describe images
+        
+        These models are completely open-access and don't require any special authentication.
+        """)
 
 # Launch the app
 demo.launch()