Latest News

SustainML integrates Retrieval-Augmented Generation (RAG) to intelligently recommend machine learning models to users — without retraining or fine-tuning large language models (LLMs). Instead, SustainML combines vector search, graph-based reasoning, and language generation to query a massive repository of models hosted on Hugging Face.

This article explains how the RAG pipeline in SustainML works under the hood, highlighting the key components: Annoy for vector search, Neo4j for structured relationships, and Mistral (through Ollama) for natural language reasoning.

Overview of the RAG Pipeline

The main goal of this RAG system is to answer questions like:

“Which Hugging Face model is best suited for translating a text from German to Spanish?”

Instead of hard-coding responses or manually browsing the Hugging Face Hub, SustainML performs a semantic search across thousands of pre-indexed models and their metadata, then uses a language model to reason about the retrieved results and produce a concise, context-aware answer.

Here’s how each part of the system contributes to the pipeline.

1. Building the Knowledge Graph in Neo4j

At the core of SustainML’s knowledge base lies a Neo4j graph database populated with models from the Hugging Face Hub. Each model becomes a node, connected to related entities such as its problem type, library, tags, and evaluation metrics.

2. Creating Semantic Embeddings and the Annoy Index

The unstructured part of model metadata (like descriptions, tags, and model cards) is embedded using SentenceTransformers.
SustainML employs the BAAI/bge-large-en model to generate 1024-dimensional embeddings that capture the semantic meaning of each model description.

These embeddings are stored in an Annoy (Approximate Nearest Neighbors Oh Yeah) index, which allows for ultra-fast retrieval of the most semantically similar models.
Annoy is ideal for large-scale RAG systems since it balances accuracy and speed without the heavy memory footprint of traditional databases.

annoy_index = AnnoyIndex(1024, 'angular')

annoy_index.load('models_index.ann')

When a user submits a query such as “Which models are good for text summarization?”, the system encodes it into an embedding and retrieves the top matching entries from the Annoy index.

3. The Retrieval Step: Semantic Search

The function search_semantic() performs this step:

query_embedding = sentence_model.encode([query], convert_to_tensor=True).cpu().numpy()[0]

nearest_neighbors = annoy_index.get_nns_by_vector(query_embedding, top_k=200)

The retrieved models (stored in a precomputed JSON file model_metadata.json) represent the knowledge passed to the reasoning step.
A small random subset is sampled to increase diversity while maintaining contextual relevance.

4. The Generation Step: LLM Reasoning with Ollama

Once the system retrieves semantically similar models, it passes them as context to a local LLM served by Ollama, using the mistral-small model.

The generate_natural_answer() function forms a structured prompt:

Based on the retrieved knowledge:

[... model metadata ...]

Answer the following question: [user question].

The format of the Hugging Face name must be like 'openai-community/gpt2-large'.

Output only one model name in JSON format.

This ensures the response is concise, machine-readable, and directly usable by downstream applications (e.g., SustainML’s orchestration layer).
A sample response might look like:

{

    "model_name": "facebook/detr-resnet-50"

}

This model name can then be fed into SustainML pipelines for automatic download, evaluation, or deployment.

5. Why Use RAG in SustainML?

SustainML’s mission is to make AI sustainable, modular, and explainable.
RAG fits perfectly into that vision:

• Efficiency – No retraining of large models. The system simply augments a fixed LLM with dynamic external knowledge.
• Transparency – The retrieved context can be inspected and explained (metadata, tags, metrics).
• Scalability – Annoy and Neo4j scale to tens of thousands of models without GPU inference bottlenecks.
• Local operation – The use of Ollama and Mistral ensures full offline reasoning without depending on remote APIs.

In short, RAG allows SustainML to combine retrieval accuracy with natural reasoning, creating a smarter model discovery assistant.

6. The Complete Flow

Here’s a summarized view of the end-to-end process:

1. User asks a question (e.g., “Best model for object detection in aerial images?”)
2. Query embedding is generated using BAAI/bge-large-en
3. Annoy index retrieves semantically related models
4. Retrieved metadata is passed to Mistral (via Ollama)
5. LLM outputs a JSON-formatted model name
6. SustainML uses the recommendation for further workflow automation

Conclusion

By blending semantic search, graph reasoning, and lightweight local LLM inference, SustainML demonstrates a practical and sustainable implementation of Retrieval-Augmented Generation.

This architecture turns the Hugging Face model hub into a searchable, context-aware recommendation engine, enabling efficient, explainable, and reproducible AI workflows — one query at a time.