Changes nemotron model

README.md

@@ -61,7 +61,7 @@ There are two AI-powered flows: **registering a sighting** and **recording help*
 
 ```mermaid
 flowchart TD
-    A["📷 Photo from phone"] --> B["🤖 Vision AI\nLlama-3.2-11B-Vision\n(HF Serverless)"]
+    A["📷 Photo from phone"] --> B["🤖 Vision AI\nLlama-3.2-11B-Vision\n(NVIDIA NIM / HF)"]
     B -->|"is it an animal?"| C{Animal detected?}
     C -->|"No"| ERR["❌ Error message\n'No dog or cat identified'"]
     C -->|"Yes"| D["📝 Structured JSON\nNemotron Mini 4B Instruct normalizes\nspecies · breed · color · condition · marks"]
@@ -106,7 +106,7 @@ The map uses color to signal urgency: **🟢 green** = dog · **🟠 orange** =
 
 | Component | Model / Library | Where it runs |
 |---|---|---|
-| Visual identification | **Llama-3.2-11B-Vision-Instruct** (11B) via HF Serverless | remote API |
+| Visual identification | **Llama-3.2-11B-Vision-Instruct** (11B) via NVIDIA NIM (or HF Serverless) | remote API |
 | Text reasoning / JSON | **Nemotron Mini 4B Instruct** via NVIDIA NIM | remote API |
 | Semantic embedding | **all-MiniLM-L6-v2** (384 dim, 22M params) | CPU, local |
 | Animal matching | Cosine similarity (threshold 0.80) | CPU, local |
@@ -148,12 +148,12 @@ git clone https://huggingface.co/spaces/build-small-hackathon/viralata-mapper
 cd viralata-mapper
 pip install -r requirements.txt
 
-# Vision via HuggingFace Serverless (Llama-3.2-11B-Vision)
-HF_TOKEN=hf_... python app.py
-
-# Add NVIDIA NIM for Nemotron Mini 4B Instruct (structured-output normalization)
+# Recommended: NVIDIA NIM powers both vision (Llama-3.2-11B-Vision) and Nemotron Mini 4B
 NVIDIA_API_KEY=nvapi_... python app.py
 
+# Alternative: vision via HuggingFace Serverless
+HF_TOKEN=hf_... python app.py
+
 # no key — runs in fallback mode without AI identification
 python app.py
 
@@ -170,8 +170,8 @@ On first run the app seeds the database with demo animals around Brasília so th
 
 | Secret | Description |
 |---|---|
-| `HF_TOKEN` | HuggingFace token for Llama-3.2-11B-Vision via Serverless Inference |
-| `NVIDIA_API_KEY` | NVIDIA NIM key for Nemotron Mini 4B Instruct (structured-output normalization) |
+| `NVIDIA_API_KEY` | NVIDIA NIM key — powers **both** vision (Llama-3.2-11B-Vision) and Nemotron Mini 4B normalization. A single key is enough. |
+| `HF_TOKEN` | Optional — vision via HF Serverless if you don't use NVIDIA NIM |
 | `MATCH_THRESHOLD` | Optional. Similarity threshold. Default: `0.80` |
 | `HF_DATASET_ID` | Optional. Hub dataset for trace publishing (e.g. `org/dataset-name`) |
 

core/ai.py

@@ -1,21 +1,26 @@
 """
-ai.py — Visao via HuggingFace InferenceClient + normalizacao via NVIDIA NIM.
+ai.py — Visao + normalizacao de texto, ambas via NVIDIA NIM (ou HF como alternativa).
 
 Arquitetura (todos os modelos < 32B):
-  - Visao  : meta-llama/Llama-3.2-11B-Vision-Instruct  (11B, HF Serverless)  -> precisa HF_TOKEN
-  - Texto  : nvidia/nemotron-mini-4b-instruct                  (4B,  NVIDIA NIM)     -> precisa NVIDIA_API_KEY (opcional)
-  - Embed  : sentence-transformers/all-MiniLM-L6-v2     (22M, local/CPU)
+  - Visao  : meta/llama-3.2-11b-vision-instruct   (11B, NVIDIA NIM)  -> NVIDIA_API_KEY
+             (ou meta-llama/Llama-3.2-11B-Vision-Instruct via HF, se so houver HF_TOKEN)
+  - Texto  : nvidia/nemotron-mini-4b-instruct      (4B,  NVIDIA NIM)  -> NVIDIA_API_KEY
+  - Embed  : sentence-transformers/all-MiniLM-L6-v2 (22M, local/CPU)
 
-O Llama enxerga a foto e devolve a descricao bruta. Se NVIDIA_API_KEY estiver
-configurada, o Nemotron Mini 4B Instruct (texto puro) normaliza esse JSON para o schema
+Fluxo: o Llama-3.2-11B-Vision enxerga a foto e devolve a descricao bruta. Em
+seguida o Nemotron Mini 4B Instruct (texto puro) normaliza esse JSON para o schema
 estrito (raca da lista permitida, cores simples, condicao, frase de matching).
-Sem NVIDIA_API_KEY, usa-se direto a saida do Llama.
+
+Basta a NVIDIA_API_KEY: ela serve tanto para a visao quanto para o Nemotron.
+Se nao houver NVIDIA_API_KEY mas houver HF_TOKEN, a visao usa o HF Serverless
+(sem a etapa de normalizacao do Nemotron).
 
 Variaveis de ambiente:
-  HF_TOKEN        — token HuggingFace (HF Serverless Inference) — obrigatorio p/ visao
-  NVIDIA_API_KEY  — chave NVIDIA NIM (free tier em build.nvidia.com) — opcional
-  HF_VISION_MODEL — override do modelo de visao (default Llama-3.2-11B-Vision)
-  NVIDIA_MODEL    — override do modelo NIM (default nvidia/nemotron-mini-4b-instruct)
+  NVIDIA_API_KEY      — chave NVIDIA NIM (free tier em build.nvidia.com)
+  HF_TOKEN            — token HuggingFace (alternativa de visao via HF Serverless)
+  NVIDIA_VISION_MODEL — override do modelo de visao NIM (default meta/llama-3.2-11b-vision-instruct)
+  NVIDIA_MODEL        — override do modelo de texto NIM (default nvidia/nemotron-mini-4b-instruct)
+  HF_VISION_MODEL     — override do modelo de visao HF (default meta-llama/Llama-3.2-11B-Vision-Instruct)
 """
 import base64
 import io
@@ -28,10 +33,9 @@ import numpy as np
 
 log = logging.getLogger(__name__)
 
-_HF_MODEL  = "meta-llama/Llama-3.2-11B-Vision-Instruct"
-# String do modelo na API NVIDIA NIM (integrate.api.nvidia.com).
-# Confirme o id exato em build.nvidia.com; pode sobrescrever via NVIDIA_MODEL.
-_NIM_MODEL = "nvidia/nemotron-mini-4b-instruct"  # Nemotron Mini 4B Instruct (texto, 4B)
+_HF_MODEL         = "meta-llama/Llama-3.2-11B-Vision-Instruct"  # id HF Serverless
+_NIM_VISION_MODEL = "meta/llama-3.2-11b-vision-instruct"        # id NVIDIA NIM (11B, free)
+_NIM_TEXT_MODEL   = "nvidia/nemotron-mini-4b-instruct"          # id NVIDIA NIM (4B, free)
 
 PROMPT = (
     "Examine this image carefully.\n"
@@ -76,38 +80,45 @@ NORMALIZE_PROMPT = (
 
 class AnimalAI:
     def __init__(self):
-        self.vision_model  = _HF_MODEL
-        self.vision_client = None  # InferenceClient (HF) — visao
-        self.nim_model     = _NIM_MODEL
-        self.nim_client    = None  # OpenAI client -> NVIDIA NIM — texto
+        self.vision_mode   = None   # "nim" | "hf"
+        self.vision_model  = None
+        self.vision_client = None   # OpenAI (NIM) ou InferenceClient (HF)
+        self.nim_text_model = _NIM_TEXT_MODEL
+        self.nim_client     = None  # OpenAI -> NVIDIA NIM, p/ normalizacao Nemotron
 
-        hf_token   = os.environ.get("HF_TOKEN", "")
         nvidia_key = os.environ.get("NVIDIA_API_KEY", "")
+        hf_token   = os.environ.get("HF_TOKEN", "")
 
-        # Visao: Llama-3.2-11B-Vision via HF Serverless (obrigatorio para identificar foto)
-        if hf_token:
-            try:
-                from huggingface_hub import InferenceClient
-                self.vision_model  = os.environ.get("HF_VISION_MODEL", _HF_MODEL)
-                self.vision_client = InferenceClient(model=self.vision_model, token=hf_token)
-                log.info("Visao: %s via HF InferenceClient", self.vision_model)
-            except ImportError:
-                log.warning("huggingface_hub nao instalado")
-        else:
-            log.warning("Sem HF_TOKEN — visao desabilitada. Configure HF_TOKEN nos Secrets.")
-
-        # Texto: Nemotron Mini 4B Instruct via NVIDIA NIM (opcional — normaliza o JSON)
+        # Cliente NVIDIA NIM (serve para visao Llama E texto Nemotron)
         if nvidia_key:
             try:
                 from openai import OpenAI
-                self.nim_model  = os.environ.get("NVIDIA_MODEL", _NIM_MODEL)
                 self.nim_client = OpenAI(
                     base_url="https://integrate.api.nvidia.com/v1",
                     api_key=nvidia_key,
                 )
-                log.info("Texto: Nemotron Mini 4B Instruct via NVIDIA NIM (%s)", self.nim_model)
+                self.nim_text_model = os.environ.get("NVIDIA_MODEL", _NIM_TEXT_MODEL)
+                # Visao via NIM (Llama-3.2-11B-Vision)
+                self.vision_mode   = "nim"
+                self.vision_model  = os.environ.get("NVIDIA_VISION_MODEL", _NIM_VISION_MODEL)
+                self.vision_client = self.nim_client
+                log.info("Visao: %s + Texto: %s via NVIDIA NIM", self.vision_model, self.nim_text_model)
+            except ImportError:
+                log.warning("openai nao instalado")
+
+        # Alternativa: visao via HF Serverless (Llama-3.2-11B-Vision) se nao houver NVIDIA
+        if self.vision_client is None and hf_token:
+            try:
+                from huggingface_hub import InferenceClient
+                self.vision_mode   = "hf"
+                self.vision_model  = os.environ.get("HF_VISION_MODEL", _HF_MODEL)
+                self.vision_client = InferenceClient(model=self.vision_model, token=hf_token)
+                log.info("Visao: %s via HF InferenceClient", self.vision_model)
             except ImportError:
-                log.warning("openai nao instalado — normalizacao Nemotron desabilitada")
+                log.warning("huggingface_hub nao instalado")
+
+        if self.vision_client is None:
+            log.warning("Sem NVIDIA_API_KEY nem HF_TOKEN — IA desabilitada. Configure um dos Secrets.")
 
         self.embedder = None
         try:
@@ -124,20 +135,29 @@ class AnimalAI:
 
         try:
             img_b64 = self._to_b64(image)
+            content = [
+                {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + img_b64}},
+                {"type": "text", "text": PROMPT},
+            ]
+
+            if self.vision_mode == "nim":
+                # NVIDIA NIM — Llama-3.2-11B-Vision (OpenAI-compatible)
+                resp = self.vision_client.chat.completions.create(
+                    model=self.vision_model,
+                    messages=[{"role": "user", "content": content}],
+                    max_tokens=400,
+                    temperature=0.1,
+                )
+                raw = resp.choices[0].message.content or ""
+            else:
+                # HuggingFace InferenceClient — suporte nativo a multimodal
+                resp = self.vision_client.chat_completion(
+                    messages=[{"role": "user", "content": content}],
+                    max_tokens=400,
+                    temperature=0.1,
+                )
+                raw = resp.choices[0].message.content or ""
 
-            # Visao multimodal — Llama-3.2-11B-Vision (HF InferenceClient)
-            resp = self.vision_client.chat_completion(
-                messages=[{
-                    "role": "user",
-                    "content": [
-                        {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + img_b64}},
-                        {"type": "text", "text": PROMPT},
-                    ],
-                }],
-                max_tokens=400,
-                temperature=0.1,
-            )
-            raw = resp.choices[0].message.content or ""
             log.info("Visao resposta: %s", raw[:200])
             result = self._parse(raw)
 
@@ -163,7 +183,7 @@ class AnimalAI:
         try:
             payload = json.dumps(result, ensure_ascii=False)
             resp = self.nim_client.chat.completions.create(
-                model=self.nim_model,
+                model=self.nim_text_model,
                 messages=[{"role": "user", "content": NORMALIZE_PROMPT + payload}],
                 max_tokens=400,
                 temperature=0.2,
@@ -178,7 +198,7 @@ class AnimalAI:
                           "description_text"):
                     if k in cleaned and cleaned[k] not in (None, "", []):
                         result[k] = cleaned[k]
-                log.info("Nemotron 4B normalizou os campos.")
+                log.info("Nemotron Mini 4B normalizou os campos.")
             return result
         except Exception as e:
             log.warning("Nemotron normalize falhou, usando saida da visao: %s", e)