DEV Community

Avani
Avani

Posted on

Running Hugging Face Inference with Kiro: From Prompt to Working Summarizer

#python #nlp #tutorial #ai

Pre-trained models are the backbone of modern NLP. Whether you're summarizing documents, classifying support tickets, or building semantic search, the Hugging Face transformers library gives you access to thousands of models, but wiring up tokenization, device management, and batching correctly takes careful attention to detail.

This post shows how Kiro handles that for you. We'll walk through a real example: building a text summarizer using a pre-trained BART model, then iterating with Kiro to add GPU support and batching, all through natural conversation.

Everything here was tested locally and verified. The output you see below is real — captured from actual script execution on a MacBook with no GPU.

The Task: Summarize Documents with a Pre-Trained Model

I wanted a script that takes paragraphs of text and produces concise summaries. The goal: use the Hugging Face transformers library directly (not the pipeline shortcut) so we have full control over tokenization, device placement, and generation parameters.

Here's what I asked Kiro:

"Write a Python script that loads a summarization model from Hugging Face and generates summaries for a list of input texts. Use the transformers library directly, not the pipeline abstraction. Pick the best model for the task."

What Kiro Generated: The Base Script

Kiro selected sshleifer/distilbart-cnn-12-6 — a distilled version of BART trained on CNN/DailyMail. It has fewer parameters (~306M vs ~406M) with a halved decoder (6 layers instead of 12), so it downloads faster and runs faster on CPU while still producing quality summaries. It used AutoModelForSeq2SeqLM (the correct class for encoder-decoder summarization models) without being told.

Here's the core of what it produced:

import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM


def load_model(model_name: str = "sshleifer/distilbart-cnn-12-6"):
    """Load the tokenizer and model from Hugging Face Hub."""
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    model.eval()  # Evaluation mode — disables dropout
    return tokenizer, model


def summarize(texts: list[str], tokenizer, model) -> list[str]:
    """Generate summaries for a list of input texts."""
    # Tokenize inputs with padding and truncation
    inputs = tokenizer(
        texts,
        padding=True,
        truncation=True,
        max_length=1024,
        return_tensors="pt",
    )

    # Generate summaries without computing gradients
    with torch.no_grad():
        summary_ids = model.generate(
            inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            max_length=130,
            min_length=30,
            num_beams=4,
            length_penalty=2.0,
            early_stopping=True,
        )

    # Decode generated token IDs back to text
    summaries = tokenizer.batch_decode(summary_ids, skip_special_tokens=True)
    return summaries
Enter fullscreen mode Exit fullscreen mode

Key Steps Explained

1. Model Loading 

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
model.eval()
Enter fullscreen mode Exit fullscreen mode

Kiro uses the Auto classes — they detect the correct architecture from the model name. For summarization, that's AutoModelForSeq2SeqLM (encoder-decoder). Under the hood, this loads a BartForConditionalGeneration model and a RobertaTokenizer. (BART shares its byte-level BPE vocabulary with RoBERTa — in fact BartTokenizer is an alias for RobertaTokenizer in the transformers library, so AutoTokenizer resolves directly to RobertaTokenizer.) model.eval() switches off dropout for consistent inference output.

2. Tokenization 

inputs = tokenizer(
    texts,
    padding=True,
    truncation=True,
    max_length=1024,
    return_tensors="pt",
)
Enter fullscreen mode Exit fullscreen mode

The tokenizer converts raw strings into tensors. padding=True ensures all sequences in a batch have equal length. truncation=True clips inputs exceeding the 1024-token context window (verified: our long-text test confirmed truncation works gracefully). return_tensors="pt" gives PyTorch tensors directly.

3. Inference (Generation) 

with torch.no_grad():
    summary_ids = model.generate(
        inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        max_length=130,
        min_length=30,
        num_beams=4,
        length_penalty=2.0,
        early_stopping=True,
    )
Enter fullscreen mode Exit fullscreen mode

Unlike classification (a single forward pass), summarization generates new tokens autoregressively. torch.no_grad() disables gradient computation for speed and memory savings. The generation parameters:

  • num_beams=4 — beam search explores 4 candidate sequences in parallel
  • length_penalty=2.0 — encourages longer, more complete summaries
  • early_stopping=True — stops when all beams produce an end token

Iterating with Kiro: Adding GPU Support

Next prompt:

"Update this to automatically use GPU if available, and fall back to CPU."

Kiro added device detection and placement:

def load_model(model_name: str = "sshleifer/distilbart-cnn-12-6"):
    """Load model with automatic device selection."""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    model.to(device)
    model.eval()

    return tokenizer, model, device


def summarize(texts: list[str], tokenizer, model, device, ...) -> list[str]:
    inputs = tokenizer(
        texts, padding=True, truncation=True, max_length=1024, return_tensors="pt",
    ).to(device)  # Move inputs to same device as model
    ...
Enter fullscreen mode Exit fullscreen mode

Tested result: On my MacBook (Apple M-series, no CUDA), the script correctly detects CUDA available: False and falls back to CPU. Apple's MPS backend is available (torch.backends.mps.is_available() = True), but Kiro correctly uses the standard CUDA/CPU pattern that works universally.

Iterating with Kiro: Adding Batching

Third prompt:

"Add batching support so I can process large datasets without running out of memory. Use a configurable batch size."

Kiro added:

def summarize_batched(
    texts: list[str],
    tokenizer,
    model,
    device,
    batch_size: int = 4,
    **kwargs,
) -> list[str]:
    """Generate summaries in batches to handle large datasets."""
    all_summaries = []

    for i in range(0, len(texts), batch_size):
        batch_texts = texts[i : i + batch_size]
        batch_summaries = summarize(
            batch_texts, tokenizer, model, device, **kwargs
        )
        all_summaries.extend(batch_summaries)

    return all_summaries
Enter fullscreen mode Exit fullscreen mode

Tested result: I ran the same 4 texts with batch_size=1, batch_size=2, and batch_size=4. All three produced identical summaries and correct result counts. The batching logic correctly slices, processes, and reassembles without losing or reordering results.

Real Output: The Final Script in Action

Here's the actual terminal output from running the tested script:

$ python summarize.py
Loading model...
Model loaded on: cpu
Generating summaries...

--- Text 1 ---
Input:   The Amazon rainforest, often referred to as the lungs of the Earth, produces approximately 20 percen...
Summary: The Amazon rainforest produces approximately 20 percent of the world's oxygen.
Spanning across nine countries in South America, it is home to an estimated 400 billion
individual trees representing over 16,000 species. Scientists continue to discover new
species every year.

--- Text 2 ---
Input:   Kubernetes has become the de facto standard for container orchestration in cloud-native applications...
Summary: Kubernetes is de facto standard for container orchestration in cloud-native
applications. It automates deployment, scaling, and management of containerized workloads
across clusters of machines. Key features include self-healing capabilities, horizontal
scaling, service discovery, and load balancing.

--- Text 3 ---
Input:   The new wireless earbuds feature active noise cancellation with three adjustable levels, 24-hour bat...
Summary: Wireless earbuds feature active noise cancellation, 24-hour battery life and IPX5
water resistance. Support Bluetooth 5.3 for stable connectivity and include touch controls
for playback and calls. Available in four colors.
Enter fullscreen mode Exit fullscreen mode

Performance: What to Expect

Measured on an Apple M-series MacBook (CPU inference):

  • Model load (from cache): 3.6s
  • Inference (3 texts, single batch): 6.5s
  • Per-text average: 2.2s
  • Total end-to-end: ~10s
  • Batched (3 texts, batch_size=2): 8.1s
  • Peak memory: ~2.5GB

First run downloads the 1.2GB model (~2 minutes on broadband). Subsequent runs load from ~/.cache/huggingface/ in seconds. On a CUDA GPU, expect 5-10x inference speedup.

Batched inference with batch_size=2 is slightly slower than processing all 3 at once (8.1s vs 6.5s) because of the overhead of two separate forward passes. The benefit shows with larger datasets where processing everything at once would exhaust memory.

What Kiro Handles Well Here

  1. Model selection — Chose distilbart-cnn-12-6 over the larger bart-large-cnn. Recognized that a demo benefits from faster downloads and inference. Used the correct AutoModelForSeq2SeqLM class without being told the architecture type.

  2. Correct API patterns — model.eval(), torch.no_grad(), proper attention_mask passing, beam search parameters, device placement. These patterns are scattered across Hugging Face docs — Kiro assembles them correctly in one pass.

  3. Iterative refinement — Each follow-up produced targeted changes. GPU support threaded a device parameter through existing functions. Batching added a wrapper reusing the existing summarize() function. No rewrites, no broken logic between iterations.

  4. Production structure — Type hints, docstrings, configurable parameters, clean function separation. The code reads like something a senior engineer would write for a shared codebase.

What I Verified Through Testing

  • Basic execution (3 sample texts): PASS — coherent summaries
  • Batch consistency (bs=1 vs bs=2 vs bs=4): PASS — identical results
  • GPU detection and fallback: PASS — graceful CPU fallback
  • Long text truncation (>1024 tokens): PASS — truncated without error
  • Short text handling: PASS — produces output (though quality degrades)
  • API signatures and return types: PASS — proper typing
  • Performance on CPU: 2.2s/text — acceptable for non-realtime use

Tips for Working with Kiro on Hugging Face Models

  • Name the task explicitly. "Summarization" tells Kiro to use AutoModelForSeq2SeqLM and model.generate(). "Classification" triggers AutoModelForSequenceClassification with softmax. The task name drives the architecture.

  • Ask for the raw API when you need control. The pipeline("summarization") abstraction hides device management, generation parameters, and batching. The direct API gives you code you can tune.

  • Let Kiro pick the model. It chose a distilled variant that's smaller and faster without sacrificing quality. You can override if you have a specific model in mind.

  • Iterate on non-functional requirements separately. Get inference working first, then add GPU, batching, or error handling. Keeps each change focused.

Run It Yourself 

pip install transformers torch
python summarize.py
Enter fullscreen mode Exit fullscreen mode

The complete script handles: model loading with automatic device selection, tokenization with padding and truncation, batched inference with configurable chunk size, beam search generation with quality parameters, and clean output formatting.

Top comments (0)