NuExtract-2.0-2B-TKG-Triples V1.1

v1.1 - NOTE: There was a minor glitch in the original upload that caused a Sampler: 3 error due to a misconfiguration in generation_config.json. It's been fixed now. Please report any problems and enjoy

NuExtract-2.0-2B-TKG-Triples is a highly specialized, edge-optimized 2B parameter model built for real-time Temporal Knowledge Graph (TKG) extraction.

Fine-tuned on a heavily refined subset of multi-session chat data, this model is designed to run locally on mobile hardware to extract precise Subject-Predicate-Object relational triples from conversational context, effectively powering on-device Long-Term Memory (LTM) architectures.

🛠 Model Specifications

• Base Architecture: Built on Qwen2.1-VL. Note: The vision tower has been bypassed during compilation; this is strictly a text-to-text execution graph.
• Parameter Count: ~2 Billion
• Export Format: .litertlm (Compatible with Google AI Edge and the LiteRT-LM framework)
• Quantization Recipe: dynamic_wi4_afp32 (Dynamic weight-only 4-bit INT4 quantization, maintaining FP32 precision on activations for accurate JSON/triple generation)
• Target Hardware: Mobile GPUs and Edge NPUs
• Footprint: ~750 MB (Ultra-lightweight for low memory pressure on edge devices)

📊 Dataset & Training

The training pipeline was strictly focused on zero-shot relational triple extraction from conversational logs:

• Base Data: william94/multi-session-chat
• Refinement: The base data was heavily filtered and refined using strict output prompting to isolate specific memory states and temporal markers.
• Sample Size: 1,500 high-quality, targeted samples.
• Fine-tuning: LoRA trained for 2 epochs (~180 steps).
• Merging: LoRA adapters merged back to base weights at FP16 prior to LiteRT compilation.

🧠 The Use Case: Temporal Knowledge Graphs & Long Term Memory

This model was created to fill a highly specific architectural role for an on-device Android project. If it proves useful for your own local AI pipelines, that's excellent. If not, it perfectly serves its intended purpose here!

Why this model?

When developing on-device Long-Term Memory (LTM), the system needs to process prompts at the end of a chat turn to extract memories. Doing this locally requires a lightweight, highly capable model.

• Gemini Nano is currently too rigidly fenced and underpowered for strict JSON data extraction.
• Gemma 4 E2B / E4B models, while excellent, are slightly too heavy to constantly run in the background just for LTM parsing without draining device resources.

The Objective

This model analyzes a conversation prompt and silently extracts relational triples—Subject, Predicate, Object clusters.

Critically, it handles temporal decay and contextual awareness. For example:

"John - was - sleeping"

When did this happen? It might be relevant context today, but in four days, a general inquiry about John shouldn't surface the fact that he was sleeping last week. The extracted triples are formulated so the parent application can attach temporal awareness to the Knowledge Graph, ensuring the AI agent only retrieves contextually relevant memories.

🚀 Deployment & Usage

Because this model is compiled into the .litertlm format, it is explicitly designed to be dropped into Google AI Edge environments using the LiteRT-LM orchestration layer.

By using the dynamic_wi4_afp32 quantization recipe, the 4-bit weights drastically reduce the memory footprint (~750MB), while the FP32 activations ensure the model doesn't suffer from precision collapse when trying to output strict Knowledge Graph brackets.

You can load and run this via the LiteRT-LM CLI or native Android Jetpack Compose pipelines.

# Example LiteRT-LM CLI execution
litert-lm run \
  --model_path=NuExtract-2.0-2B-TKG-Triples.litertlm \
  --backend=gpu \
  --prompt="[Chat log here for extraction]"