Spaces:
Running on Zero
Running on Zero
| import spaces # before torch, so ZeroGPU can manage CUDA | |
| import gradio as gr | |
| import torch | |
| from peft import PeftModel | |
| from PIL import Image | |
| from transformers import AutoProcessor, AutoTokenizer, Qwen3_5ForConditionalGeneration | |
| from jemm import DecisionModel | |
| from jemm.model import _decision_config | |
| BASE = "Qwen/Qwen3.8-27B" | |
| ADAPTER = "MaestroYan/JEMM" | |
| MAX_PIXELS = 1280 * 800 | |
| # Attach the adapter on CPU, then .to("cuda") once, so ZeroGPU can pack all weights at startup | |
| base = Qwen3_5ForConditionalGeneration.from_pretrained(BASE, dtype=torch.bfloat16, attn_implementation="sdpa") | |
| peft_model = PeftModel.from_pretrained(base, ADAPTER, is_trainable=False, torch_device="cpu").to("cuda") | |
| model = DecisionModel(peft_model, AutoTokenizer.from_pretrained(BASE), name="JEMM", | |
| processor=AutoProcessor.from_pretrained(BASE), **_decision_config(ADAPTER)) | |
| def fit(image: Image.Image) -> Image.Image: | |
| image = image.convert("RGB") | |
| scale = (MAX_PIXELS / (image.width * image.height)) ** 0.5 | |
| if scale >= 1: | |
| return image | |
| return image.resize((round(image.width * scale), round(image.height * scale)), Image.Resampling.LANCZOS) | |
| def decide(state: str, question: str, candidates: str, image: Image.Image | None = None): | |
| """Pick one candidate for the question and return a probability for every candidate. | |
| Args: | |
| state: Context the decision is based on, such as a conversation or a task description. | |
| question: The question to answer. | |
| candidates: 2 to 32 distinct candidates, one per line. | |
| image: Optional screenshot. Larger images are scaled to about 1280x800, the training size. | |
| """ | |
| options = [line.strip() for line in candidates.splitlines() if line.strip()] | |
| if not 2 <= len(options) <= 32 or len(set(options)) != len(options): | |
| raise gr.Error("Enter 2 to 32 distinct candidates, one per line.") | |
| questions = {"q": {"type": "choice", "instructions": question, "criteria": {o: o for o in options}}} | |
| try: | |
| result = model.decide(state, questions, None if image is None else [fit(image)])["q"] | |
| except ValueError as exc: | |
| raise gr.Error(str(exc)) from exc | |
| verdict = "" if result["confidence"] >= model.threshold else " · below threshold, treat as undecided" | |
| return result["probabilities"], f"**{result['choice']}** \nConfidence {result['confidence']:.3f}{verdict}" | |
| TEXT_EXAMPLES = [ | |
| ["User: Will it rain in Paris tomorrow?", "Which tool should handle the request?", | |
| "get_weather: weather forecast for a city\nsearch_flights: flight search\nNo tool applies"], | |
| ["Customer: I was charged twice for my March invoice. Please refund one of the charges.", | |
| "Which team should handle this ticket?", "Billing\nTechnical support\nSales\nAccount security"], | |
| ] | |
| SCREENSHOT_EXAMPLES = [ | |
| ["Task: Search for wireless headphones.\nPrevious actions: [textbox] Search -> TYPE: wireless headphones", | |
| "Which element should be clicked next? Candidates refer to the elements marked with the same letters on the screenshot.", | |
| "\n".join(f"CLICK the element marked {x}" for x in "ABCDE"), "examples/search.png"], | |
| ] | |
| HEADER = """# JEMM | |
| **Like [Jev](https://typesafe.ai/blog/introducing-system-one-models-and-jev), but multimodal and open-weight.** Picks one candidate per question and returns a probability for every candidate, from text or text plus a screenshot. | |
| [Model card](https://huggingface.co/MaestroYan/JEMM) · [GitHub](https://github.com/ypcypc/JEMM) · Not affiliated with TypeSafe AI.""" | |
| def panel(name: str, examples: list, image: bool): | |
| with gr.Tab(name): | |
| with gr.Row(equal_height=False): | |
| with gr.Column(): | |
| state = gr.Textbox(label="State", lines=4) | |
| question = gr.Textbox(label="Question") | |
| candidates = gr.Textbox(label="Candidates, one per line", lines=5) | |
| inputs = [state, question, candidates] | |
| if image: | |
| inputs.append(gr.Image(label="Screenshot", type="pil")) | |
| button = gr.Button("Decide", variant="primary") | |
| with gr.Column(): | |
| result = gr.Markdown() | |
| probabilities = gr.Label(label="Probabilities", show_label=True) | |
| button.click(decide, inputs, [probabilities, result], api_name="decide" if not image else "decide_screenshot") | |
| gr.Examples(examples, inputs, cache_examples=False) | |
| with gr.Blocks(title="JEMM") as demo: | |
| gr.Markdown(HEADER) | |
| panel("Text", TEXT_EXAMPLES, image=False) | |
| panel("Screenshot", SCREENSHOT_EXAMPLES, image=True) | |
| demo.launch(theme=gr.themes.Default(primary_hue="blue", font=[gr.themes.GoogleFont("Inter"), "sans-serif"])) | |