app.py

"""
Goals: Off-Brand, Llama Champion, Field Notes

The plan
Create an LLM powered transpiler
Taking input code (Python, JS, etc) and translate it into Mermaid.js syntax
Gradio then interprets that visually.

The Pipeline:
1. Input. User pastes code into a gradio.Code() block.
2. Process. Send that code to Small Model with a specific system prompt
3. Graph. Capture the resulting mermaid string and visualize it

- include few shot examples in system prompt
- master prompting for the system prompt and test differnet ones

Future ideas include allowing user to edit the resulting flowchart and thereby edit code

Model:
https://huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct-GGUF

To do
- need to create the basic gradio looks & pipeline

"""
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import gradio as gr
from typing import Any, cast # to resolve PyLance freaking out over llama-cpp-python in the generate_flowchart function

# ----- Get Model ----- #
# Download Q4_K_M GGUF file from the repo
model_path = hf_hub_download(
    repo_id="Qwen/Qwen2.5-Coder-7B-Instruct-GGUF",
    filename="qwen2.5-coder-7b-instruct-q4_k_m.gguf"
)

# Initialize llama.cpp with the local cached path
llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=2
)

# ----- Generation function ----- #
def generate_flowchart(src_code: str):
    # check if src_code is empty
    if not src_code.strip(): return ""

    system_prompt = (
        "..."
        "..."
    )

    # Casting else PyLance gets mad
    response = cast(Any, llm.create_chat_completion(
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": src_code}
        ],
        temperature=0.1, # Keep it quite deterministic for now
        max_tokens=1024,
        stream=False
    ))

    content = response["choices"][0]["message"]["content"]

    # Fallback to fix the .strip() on None error
    mermaid_raw = (content or "").strip()

    return mermaid_raw

# ----- Gradio Interface (Basic, Not Custom, Archive Later) ----- #

# ----- Gradio Interface (Custom) ----- #