Update app.py

app.py

@@ -1,93 +1,147 @@
 import gradio as gr
-from awq import AutoAWQForCausalLM
-from transformers import AutoTokenizer, AutoConfig
-from huggingface_hub import HfApi, login
+import pandas as pd
+import matplotlib.pyplot as plt
+import io
+import google.generativeai as genai
+from PIL import Image
+import ast
+import re
+import traceback
 
-def quantize_model(
-    model_id: str,
-    hf_token: str,
-    repo_name: str,
-    progress=gr.Progress(track_tqdm=True)
-):
+def process_file(api_key, file, instructions):
+    # Configure Gemini API with error handling
     try:
-        # Validate credentials first
-        login(token=hf_token, add_to_git_credential=True)
-        api = HfApi(token=hf_token)
-        
-        # Check model accessibility
+        genai.configure(api_key=api_key)
+        model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')
+    except Exception as e:
+        return [generate_error_image(f"API Config Error: {str(e)}")] * 3
+
+    # Load data with validation
+    try:
+        if file.name.endswith('.csv'):
+            df = pd.read_csv(file.name)
+        else:
+            df = pd.read_excel(file.name)
+            
+        if df.empty:
+            raise ValueError("Empty dataset uploaded")
+        if len(df.columns) < 2:
+            raise ValueError("Dataset needs at least 2 columns")
+    except Exception as e:
+        return [generate_error_image(f"Data Error: {str(e)}")] * 3
+
+    # Enhanced prompt with strict formatting
+    prompt = f"""Generate 3 Python matplotlib codes with these rules:
+    1. Perfect Python 3.10 syntax with 4-space indentation
+    2. Complete code blocks with proper indentation
+    3. Use ONLY these variables: df, plt
+    4. Each visualization must:
+       - Start with: plt.figure(figsize=(16,9), dpi=120)
+       - Use plt.style.use('ggplot')
+       - Include title, axis labels, and data visualization
+       - End with plt.tight_layout()
+    
+    Dataset Columns: {list(df.columns)}
+    Data Sample: {df.head(3).to_dict()}
+    User Instructions: {instructions or 'Show general trends'}
+    
+    Format EXACTLY as:
+    # Visualization 1
+    [properly indented code]
+    
+    # Visualization 2
+    [properly indented code]
+    
+    # Visualization 3
+    [properly indented code]
+    """
+
+    # Get and process Gemini response
+    try:
+        response = model.generate_content(prompt)
+        code_blocks = re.split(r'# Visualization \d+\s*', response.text)[1:4]
+    except Exception as e:
+        return [generate_error_image("API Response Error")] * 3
+
+    visualizations = []
+    for i, block in enumerate(code_blocks, 1):
         try:
-            api.model_info(model_id)
+            # Clean and validate code
+            cleaned_code = sanitize_code(block, df.columns)
+            img = execute_plot_code(cleaned_code, df)
+            visualizations.append(img)
         except Exception as e:
-            raise ValueError(f"Model access error: {str(e)}. Check:\n1. Token permissions\n2. Model existence\n3. Accept model terms at https://huggingface.co/{model_id}")
+            print(f"Visualization {i} Error:\n{traceback.format_exc()}")
+            visualizations.append(generate_error_image(f"Error in Viz {i}"))
 
-        # Load config with proper auth
-        config = AutoConfig.from_pretrained(
-            model_id,
-            token=hf_token,
-            trust_remote_code=True
-        )
-        
-        # Handle Llama 3 rope_scaling
-        if hasattr(config, 'rope_scaling') and isinstance(config.rope_scaling, dict):
-            config.rope_scaling = {
-                "type": config.rope_scaling.get("rope_type", "linear"),
-                "factor": config.rope_scaling.get("factor", 1.0)
-            }
-        
-        # Load model with validated credentials
-        model = AutoAWQForCausalLM.from_pretrained(
-            model_id,
-            config=config,
-            token=hf_token,
-            trust_remote_code=True,
-            device_map="auto"
-        )
+    # Ensure exactly 3 outputs
+    return visualizations + [generate_error_image("Not Generated")]*(3-len(visualizations))
+
+def sanitize_code(code_block, columns):
+    """Fix indentation and syntax issues in generated code"""
+    lines = []
+    indent_level = 0
+    indent_size = 4
+    stack = []
+
+    for line in code_block.split('\n'):
+        line = line.rstrip()
+        if not line:
+            continue
+
+        # Remove markdown artifacts
+        line = re.sub(r'^```python|```$', '', line)
         
-        # Load tokenizer with same credentials
-        tokenizer = AutoTokenizer.from_pretrained(
-            model_id,
-            token=hf_token,
-            trust_remote_code=True
-        )
+        # Handle indentation triggers
+        if re.match(r'^\s*(for|if|while|with|def|class|try|except|else|elif)\b', line):
+            stack.append(indent_level)
+            indent_level += 1
+        elif re.match(r'^\s*(return|pass|break|continue|raise)', line):
+            indent_level = max(0, indent_level - 1)
         
-        # Quantize with auto-detected settings
-        model.quantize(tokenizer, quant_config={
-            "zero_point": True,
-            "q_group_size": 128,
-            "w_bit": 4,
-            "version": "GEMM" if "llama" in model_id.lower() else "GEMV"
-        })
+        # Apply current indentation
+        current_indent = ' ' * (indent_level * indent_size)
+        cleaned_line = current_indent + line.lstrip()
         
-        # Save and push
-        save_path = f"{model_id.split('/')[-1]}-awq"
-        model.save_quantized(save_path)
-        model.push_to_hub(repo_name, token=hf_token)
+        # Check for dedent patterns
+        if re.match(r'^\s*(\}|\)|]|else:|elif |except)', line):
+            if stack:
+                indent_level = stack.pop()
         
-        return f"✅ Success!\nSaved: {save_path}\nPushed to: {repo_name}"
+        lines.append(cleaned_line)
     
-    except Exception as e:
-        return f"❌ Critical Error:\n{str(e)}"
-
-with gr.Blocks() as app:
-    gr.Markdown("## 🔐 Secure AutoAWQ Quantizer")
+    cleaned_code = '\n'.join(lines)
     
-    with gr.Row():
-        model_id = gr.Textbox(label="Model ID", 
-                            placeholder="meta-llama/Meta-Llama-3-8B-Instruct",
-                            info="Must have access rights")
-        hf_token = gr.Textbox(label="HF Token", 
-                            type="password",
-                            info="Required for gated models")
-        repo_name = gr.Textbox(label="Destination Repo",
-                             info="Format: username/repo-name")
+    # Validate syntax
+    try:
+        ast.parse(cleaned_code)
+    except SyntaxError as e:
+        raise ValueError(f"Syntax Error: {str(e)}")
     
-    go_btn = gr.Button("Start Quantization", variant="primary")
-    output = gr.Markdown()
+    return cleaned_code
+
+def execute_plot_code(code, df):
+    """Safely execute plotting code with resource management"""
+    buf = io.BytesIO()
+    plt.figure(figsize=(16, 9), dpi=120)
+    plt.style.use('ggplot')
     
-    go_btn.click(
-        quantize_model,
-        inputs=[model_id, hf_token, repo_name],
-        outputs=output
-    )
+    try:
+        exec(code, {'df': df, 'plt': plt})
+        plt.tight_layout()
+        plt.savefig(buf, format='png', bbox_inches='tight')
+        buf.seek(0)
+        return Image.open(buf)
+    except Exception as e:
+        raise RuntimeError(f"Execution Error: {str(e)}")
+    finally:
+        plt.close()
 
-app.launch()
+def generate_error_image(message):
+    """Create error indication image with message"""
+    from PIL import ImageDraw, ImageFont
+    
+    img = Image.new('RGB', (1920, 1080), color=(255, 255, 255))
+    try:
+        draw = ImageDraw.Draw(img)
+        font = ImageFont