Update app.py

app.py

@@ -3,6 +3,7 @@ import torch
 import sys
 import traceback
 import os
+from huggingface_hub import login
 
 def system_info():
     try:
@@ -26,162 +27,74 @@ def system_info():
     except Exception as e:
         return f"Error: {str(e)}\n\n{traceback.format_exc()}"
 
-def test_phi3_mini():
+def test_gemma3():
     try:
         result = []
-        result.append("Testing Phi-3 Mini model...")
+        result.append("Testing Gemma 3 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
-        
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        
-        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"
+        # Get token from environment
+        token = os.environ.get("HUGGINGFACE_TOKEN", "")
+        if token:
+            result.append(f"Token found: {token[:5]}...")
+        else:
+            result.append("No token found in environment variables!")
+            return "\n".join(result)
+        
+        # Login to Hugging Face
+        try:
+            login(token=token)
+            result.append("Successfully logged in to Hugging Face Hub")
+        except Exception as e:
+            result.append(f"Error logging in: {e}")
+            return "\n".join(result)
+        
+        # Use Gemma 3 GGUF model
+        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...")
+        from huggingface_hub import hf_hub_download
+        
+        model_path = hf_hub_download(
+            repo_id=model_id,
+            filename=model_filename,
+            token=token
         )
-        
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            quantization_config=quantization_config,
-            device_map="auto"
+        result.append(f"Model downloaded to: {model_path}")
+        
+        # Load the model
+        result.append("Loading model...")
+        try:
+            import llama_cpp
+        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
+        
+        from llama_cpp import Llama
+        
+        llm = Llama(
+            model_path=model_path,
+            n_ctx=2048,  # Context window size
+            n_gpu_layers=-1  # Use all available GPU layers
         )
         
+        # Generate text
         result.append("Generating text...")
         prompt = "Write a short poem about artificial intelligence."
         
-        # 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"
+        output = llm(
+            prompt,
+            max_tokens=100,
+            temperature=0.7,
+            top_p=0.95,
+            echo=False
         )
         
-        # 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!")
+        generated_text = output["choices"][0]["text"]
+        result.append(f"Generated text: {generated_text}")
+        result.append("Gemma 3 test successful!")
         
         return "\n".join(result)
     except Exception as e:
@@ -190,7 +103,7 @@ def test_phi3_with_image():
 # Create Gradio interface
 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.")
+    gr.Markdown("Testing Gemma 3 model for StaffManager application.")
     
     with gr.Tab("System Info"):
         with gr.Row():
@@ -205,56 +118,28 @@ with gr.Blocks(title="StaffManager AI Assistant") as demo:
             outputs=[info_result]
         )
     
-    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.Tab("Gemma 3 Test"):
         with gr.Row():
             with gr.Column():
-                combined_button = gr.Button("Generate Image Description")
+                gemma_button = gr.Button("Test Gemma 3")
             with gr.Column():
-                combined_result = gr.Textbox(label="Description Results", lines=20)
+                gemma_result = gr.Textbox(label="Test Results", lines=20)
         
-        combined_button.click(
-            fn=test_phi3_with_image,
+        gemma_button.click(
+            fn=test_gemma3,
             inputs=[],
-            outputs=[combined_result]
+            outputs=[gemma_result]
         )
     
     with gr.Tab("About"):
         gr.Markdown("""
         ## About StaffManager AI Assistant
         
-        This Space demonstrates AI capabilities for StaffManager using open-access models:
+        This Space tests the Gemma 3 model for the StaffManager application.
         
-        - **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
+        - **Gemma 3**: Google's 27B parameter model in GGUF format for efficient inference
         
-        These models are completely open-access and don't require any special authentication.
+        This model requires authentication with a Hugging Face token that has been granted access to the model.
         """)
 
 # Launch the app