Spaces:
Running
Running
| import os | |
| import zipfile | |
| import shutil | |
| import gradio as gr | |
| from huggingface_hub import hf_hub_download | |
| from llama_cpp import Llama | |
| # Fast download enabled | |
| os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1" | |
| CACHE_DIR = os.path.join(os.getcwd(), "model_cache") | |
| SANDBOX_DIR = os.path.join(os.getcwd(), "sandbox") | |
| # Ensure sandbox directory exists | |
| if not os.path.exists(SANDBOX_DIR): | |
| os.makedirs(SANDBOX_DIR) | |
| llm = None | |
| def load_model(): | |
| global llm | |
| print("--- 📥 Downloading 26B Model ---") | |
| model_path = hf_hub_download( | |
| repo_id="BugTraceAI/BugTraceAI-Apex-G4-26B-Q4", | |
| filename="BugTraceAI-Apex-G4-26B-Q4.gguf", | |
| cache_dir=CACHE_DIR | |
| ) | |
| llm = Llama(model_path=model_path, n_ctx=2048, n_threads=2, use_mmap=True) | |
| return "Model Ready!" | |
| def handle_zip(file_obj): | |
| if file_obj is None: | |
| return "No file uploaded." | |
| # Clear previous sandbox data | |
| if os.path.exists(SANDBOX_DIR): | |
| shutil.rmtree(SANDBOX_DIR) | |
| os.makedirs(SANDBOX_DIR) | |
| try: | |
| if file_obj.name.endswith(".zip"): | |
| with zipfile.ZipFile(file_obj.name, 'r') as zip_ref: | |
| zip_ref.extractall(SANDBOX_DIR) | |
| # List files for the AI context | |
| files = [] | |
| for root, dirs, filenames in os.walk(SANDBOX_DIR): | |
| for f in filenames: | |
| files.append(os.path.relpath(os.path.join(root, f), SANDBOX_DIR)) | |
| file_list_str = "\n".join(files[:20]) # Limit to first 20 files to save tokens | |
| return f"ZIP Extracted. Files found:\n{file_list_str}" | |
| else: | |
| return "Please upload a .zip file." | |
| except Exception as e: | |
| return f"Error: {e}" | |
| def chat_func(message, history, zip_file): | |
| if llm is None: | |
| return "Model is still loading... wait ~5 mins." | |
| # Logic to tell the AI what is in the sandbox | |
| context = "" | |
| if zip_file: | |
| files_in_sandbox = [f for f in os.listdir(SANDBOX_DIR)] | |
| context = f"The user has uploaded a ZIP. Files available in /sandbox: {files_in_sandbox}\n" | |
| prompt = f"{context}User: {message}\nAssistant:" | |
| output = llm(prompt, max_tokens=512, stop=["User:"]) | |
| return output["choices"][0]["text"] | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# 📦 ZIP-Enabled AI Sandbox") | |
| with gr.Row(): | |
| zip_input = gr.File(label="Upload ZIP Project", file_types=[".zip"]) | |
| status_output = gr.Textbox(label="Extraction Status") | |
| zip_input.change(handle_zip, inputs=zip_input, outputs=status_output) | |
| gr.ChatInterface(fn=chat_func, additional_inputs=[zip_input]) | |
| if __name__ == "__main__": | |
| import threading | |
| threading.Thread(target=load_model).start() | |
| demo.launch(server_name="0.0.0.0", server_port=7860) |