Qwen.py

import os
from huggingface_hub import InferenceClient

def generate_mermaid_from_qwen(base64_image):
    """
    Generate Mermaid diagram using the Qwen2.5-VL-7B-Instruct model on Hugging Face Hub.
    Requires the Hyperbolic provider and an API key.
    """
    client = InferenceClient(
        provider="hyperbolic",
        api_key=os.environ.get("HF_HUGGINGFACE_TOKEN", ""),
    )

    # Mermaid prompt adapted from your Gemma version
    completion = client.chat.completions.create(
        model="Qwen/Qwen2.5-VL-7B-Instruct",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": (
                            "You must generate Mermaid code ONLY. "
                            "Start directly with 'graph TD'. "
                            "No comments, no explanations, no HTML tags like '<br>', no links, "
                            "no quotes like '\"\"'. Only pure diagram code. "
                            "Now, based on this image, generate the Mermaid diagram structure."
                        )
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}"
                        }
                    }
                ]
            }
        ],
        max_tokens=1000,
        temperature=0.8,
        top_p=0.9,
    )

    raw = completion.choices[0].message.content

    # Clean Mermaid code
    clean_lines = []
    for line in raw.split("\n"):
        line = line.split("```")[0].split("'''")[0].strip()
        if line and not line.startswith("style"):
            clean_lines.append(line)

    return "\n".join(clean_lines)