|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import base64 |
| 4 | + |
| 5 | +from celery import shared_task |
| 6 | +from django.conf import settings |
| 7 | +from django_large_image import utilities |
| 8 | +import large_image |
| 9 | + |
| 10 | +from uvdat.core.models import RasterData, TaskResult |
| 11 | + |
| 12 | +from .analysis_type import AnalysisInputError, AnalysisTask, AnalysisType |
| 13 | + |
| 14 | +MODEL_CARD_URL = "https://huggingface.co/unsloth/Qwen3.5-9B-GGUF" |
| 15 | +SYSTEM_PROMPT = ( |
| 16 | + "You are an urban planning and geospatial analysis expert specializing in " |
| 17 | + "land use patterns, hydrology, transportation networks, and municipal policy. " |
| 18 | + "Analyze the provided imagery to answer the user's question. In your answer, " |
| 19 | + "assume that the user is also a geospatial analyst with the same expertise." |
| 20 | +) |
| 21 | +TOKEN_RANGE = {"min": 1000, "max": 10000, "step": 1000} |
| 22 | +MAX_PROMPT_LENGTH = 4000 |
| 23 | +THUMBNAIL_SIZE = 4000 |
| 24 | +MAX_STARTUP_WAIT = 300 |
| 25 | + |
| 26 | + |
| 27 | +class ImageryAskQwen(AnalysisType): |
| 28 | + def __init__(self): |
| 29 | + super().__init__() |
| 30 | + self.name = "Imagery: Ask Qwen" |
| 31 | + self.description = "Select an imagery layer and ask Qwen 3.5 about it." |
| 32 | + self.details = ( |
| 33 | + "Inferencing with unsloth/Qwen3.5-9B-GGUF provided by a " |
| 34 | + "Kitware-hosted Huggingface Inference Endpoint. " |
| 35 | + f"See the model card at {MODEL_CARD_URL}. " |
| 36 | + "Responses may cut off mid-sentence if max_tokens is reached." |
| 37 | + ) |
| 38 | + self.db_value = "imagery_ask_qwen" |
| 39 | + self.input_types = { |
| 40 | + "imagery": "RasterData", |
| 41 | + "text_prompt": "string", |
| 42 | + "max_tokens": "number", |
| 43 | + } |
| 44 | + self.output_types = { |
| 45 | + "response": "markdown", |
| 46 | + } |
| 47 | + self.attribution = "Unsloth AI, Kitware Inc." |
| 48 | + |
| 49 | + @classmethod |
| 50 | + def is_enabled(cls) -> bool: |
| 51 | + return ( |
| 52 | + settings.UVDAT_ENABLE_IMAGERY_ASK_QWEN |
| 53 | + and settings.UVDAT_HF_TOKEN is not None |
| 54 | + and settings.UVDAT_HF_NAMESPACE is not None |
| 55 | + and settings.UVDAT_HF_ENDPOINT_NAMES.get("qwen") is not None |
| 56 | + ) |
| 57 | + |
| 58 | + def get_input_options(self): |
| 59 | + return { |
| 60 | + "imagery": RasterData.objects.filter(dataset__category="imagery"), |
| 61 | + "text_prompt": [], |
| 62 | + "max_tokens": [TOKEN_RANGE], |
| 63 | + } |
| 64 | + |
| 65 | + def validate_inputs(self, inputs): |
| 66 | + super().validate_inputs(inputs) |
| 67 | + try: |
| 68 | + imagery = RasterData.objects.get(id=inputs.get("imagery")) |
| 69 | + except RasterData.DoesNotExist as e: |
| 70 | + err_msg = "Imagery raster does not exist." |
| 71 | + raise AnalysisInputError(err_msg) from e |
| 72 | + if imagery.dataset.category != "imagery": |
| 73 | + err_msg = 'Selected raster is not categorized as "imagery".' |
| 74 | + raise AnalysisInputError(err_msg) |
| 75 | + text_prompt = str(inputs.get("text_prompt")) |
| 76 | + if len(text_prompt) > MAX_PROMPT_LENGTH: |
| 77 | + err_msg = f"Prompt too long. Provide a prompt with <{MAX_PROMPT_LENGTH} characters." |
| 78 | + raise AnalysisInputError(err_msg) |
| 79 | + max_tokens = int(inputs.get("max_tokens")) |
| 80 | + if max_tokens < TOKEN_RANGE["min"] or max_tokens > TOKEN_RANGE["max"]: |
| 81 | + err_msg = f"max_tokens must be between {TOKEN_RANGE['min']} and {TOKEN_RANGE['max']}." |
| 82 | + raise AnalysisInputError(err_msg) |
| 83 | + |
| 84 | + def run_task(self, *, project, **inputs): |
| 85 | + text_prompt = inputs.get("text_prompt") |
| 86 | + result = TaskResult.objects.create( |
| 87 | + name=text_prompt[:250], |
| 88 | + task_type=self.db_value, |
| 89 | + inputs=inputs, |
| 90 | + project=project, |
| 91 | + status="Initializing Task...", |
| 92 | + ) |
| 93 | + imagery_ask_qwen.delay(result.id) |
| 94 | + return result |
| 95 | + |
| 96 | + def finalize(self, result): |
| 97 | + pass |
| 98 | + |
| 99 | + |
| 100 | +@shared_task(base=AnalysisTask) |
| 101 | +def imagery_ask_qwen(result_id): |
| 102 | + # Only available with [tasks] extra |
| 103 | + from huggingface_hub import ( # noqa: PLC0415 |
| 104 | + InferenceEndpointTimeoutError, |
| 105 | + get_inference_endpoint, |
| 106 | + ) |
| 107 | + |
| 108 | + result = TaskResult.objects.get(id=result_id) |
| 109 | + imagery = RasterData.objects.get(id=result.inputs.get("imagery")) |
| 110 | + text_prompt = result.inputs.get("text_prompt") |
| 111 | + max_tokens = int(result.inputs.get("max_tokens")) |
| 112 | + |
| 113 | + result.write_status("Encoding imagery...") |
| 114 | + imagery_path = utilities.field_file_to_local_path(imagery.cloud_optimized_geotiff) |
| 115 | + src = large_image.open(imagery_path) |
| 116 | + thumbnail_bytes, _ = src.getThumbnail(THUMBNAIL_SIZE, THUMBNAIL_SIZE, encoding="PNG") |
| 117 | + thumbnail_b64 = base64.b64encode(thumbnail_bytes).decode("utf-8") |
| 118 | + thumbnail_uri = f"data:image/jpeg;base64,{thumbnail_b64}" |
| 119 | + |
| 120 | + result.write_status("Starting inference endpoint...") |
| 121 | + endpoint = get_inference_endpoint( |
| 122 | + name=settings.UVDAT_HF_ENDPOINT_NAMES["qwen"], |
| 123 | + namespace=settings.UVDAT_HF_NAMESPACE, |
| 124 | + token=settings.UVDAT_HF_TOKEN, |
| 125 | + ) |
| 126 | + endpoint.resume() |
| 127 | + try: |
| 128 | + endpoint.wait(timeout=MAX_STARTUP_WAIT) |
| 129 | + except InferenceEndpointTimeoutError: |
| 130 | + result.write_error("Endpoint failed to start in 5 minutes. Try again later.") |
| 131 | + return |
| 132 | + |
| 133 | + result.write_status("Sending question to Qwen...") |
| 134 | + messages = [ |
| 135 | + { |
| 136 | + "role": "system", |
| 137 | + "content": SYSTEM_PROMPT, |
| 138 | + }, |
| 139 | + { |
| 140 | + "role": "user", |
| 141 | + "content": [ |
| 142 | + {"type": "image_url", "image_url": {"url": thumbnail_uri}}, |
| 143 | + {"type": "text", "text": text_prompt}, |
| 144 | + ], |
| 145 | + }, |
| 146 | + ] |
| 147 | + |
| 148 | + result.write_status("Awaiting Qwen's response...") |
| 149 | + chat = endpoint.client.chat_completion( |
| 150 | + model="unsloth/Qwen3.5-9B-GGUF", |
| 151 | + messages=messages, |
| 152 | + max_tokens=max_tokens, |
| 153 | + ) |
| 154 | + response = "" |
| 155 | + for choice in chat.choices: |
| 156 | + if choice.finish_reason == "length": |
| 157 | + # max tokens exceeded, use reasoning content |
| 158 | + response += choice.message.reasoning_content |
| 159 | + else: |
| 160 | + response += choice.message.content |
| 161 | + result.write_outputs({"response": response}) |
0 commit comments