Update app.py

app.py

@@ -1,76 +1,120 @@
 import gradio as gr
 import torch
+import sys
+import traceback
+import os
+from huggingface_hub import hf_hub_download
 
-# Create a simple interface
-with gr.Blocks(title="Simple Qwen Test") as demo:
-    gr.Markdown("# Simple Qwen Test")
-    gr.Markdown("This is a minimal test to check if the Space is working.")
-    
-    with gr.Tab("Basic Test"):
-        with gr.Row():
-            with gr.Column():
-                test_button = gr.Button("Run Basic Test")
-            with gr.Column():
-                test_result = gr.Textbox(label="Test Results", lines=10)
+def system_info():
+    try:
+        import transformers
         
-        def basic_test():
-            try:
-                # Just print system info
-                import sys
-                import transformers
-                
-                result = []
-                result.append(f"Python version: {sys.version}")
-                result.append(f"PyTorch version: {torch.__version__}")
-                result.append(f"Transformers version: {transformers.__version__}")
-                result.append("Basic test successful!")
-                
-                return "\n".join(result)
-            except Exception as e:
-                import traceback
-                return f"Error: {str(e)}\n\n{traceback.format_exc()}"
+        result = []
+        result.append(f"Python version: {sys.version}")
+        result.append(f"PyTorch version: {torch.__version__}")
+        result.append(f"Transformers version: {transformers.__version__}")
         
-        test_button.click(
-            fn=basic_test,
-            inputs=[],
-            outputs=[test_result]
-        )
+        # Check GPU availability
+        if torch.cuda.is_available():
+            result.append(f"GPU available: {torch.cuda.get_device_name(0)}")
+            result.append(f"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
+        else:
+            result.append("No GPU available")
+            
+        result.append("System info test successful!")
+        
+        return "\n".join(result)
+    except Exception as e:
+        return f"Error: {str(e)}\n\n{traceback.format_exc()}"
 
-# Launch the app
-def test_qwen_text():
+def test_gemma_gguf():
     try:
-        # Use Qwen model with 4-bit quantization to reduce memory usage
-        model_id = "Qwen/Qwen2-7B-Instruct"
-        
         result = []
-        result.append("Loading tokenizer...")
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        result.append("Testing Gemma 3 GGUF model...")
         
-        result.append("Loading model with quantization...")
-        from transformers import BitsAndBytesConfig
+        # 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__}")
         
-        quantization_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_compute_dtype=torch.float16,
-            bnb_4bit_quant_type="nf4"
+        # 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}")
+        
+        # Load the model
+        result.append("Loading model...")
+        from llama_cpp import Llama
         
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            quantization_config=quantization_config,
-            device_map="auto"
+        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 AI."
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-        outputs = model.generate(**inputs, max_new_tokens=50)
-        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        prompt = "Write a short poem about artificial intelligence."
+        
+        output = llm(
+            prompt,
+            max_tokens=100,
+            temperature=0.7,
+            top_p=0.95,
+            echo=False
+        )
         
+        generated_text = output["choices"][0]["text"]
         result.append(f"Generated text: {generated_text}")
-        result.append("Qwen text model test successful!")
+        result.append("Gemma 3 GGUF 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.Tab("System Info"):
+        with gr.Row():
+            with gr.Column():
+                info_button = gr.Button("Get System Info")
+            with gr.Column():
+                info_result = gr.Textbox(label="System Information", lines=10)
+        
+        info_button.click(
+            fn=system_info,
+            inputs=[],
+            outputs=[info_result]
+        )
+    
+    with gr.Tab("Gemma 3 GGUF Test"):
+        with gr.Row():
+            with gr.Column():
+                gemma_button = gr.Button("Test Gemma 3 GGUF")
+            with gr.Column():
+                gemma_result = gr.Textbox(label="Test Results", lines=20)
+        
+        gemma_button.click(
+            fn=test_gemma_gguf,
+            inputs=[],
+            outputs=[gemma_result]
+        )
+
+# Launch the app
 demo.launch()