Back to Home

Russian-language AI for X-ray analysis: assembling ViT + ruGPT-3 model

The article describes practical experience in creating a multimodal neural network for analyzing X-ray images and generating medical reports in Russian. It covers the stages of combining ViT and ruGPT-3 models, data preparation, training, and deploying the solution.

Assembled a Russian-language AI radiologist: from architecture to inference
Advertisement 728x90

Building a Russian Multimodal AI Radiologist: ViT + ruGPT-3 Kaggle Guide

Developing medical AI solutions in Russian faces a shortage of ready-to-use tools and datasets. This article walks through a hands-on case study of assembling and training a VisionEncoderDecoder model to generate medical reports from X-rays, combining the ViT visual encoder with the ruGPT-3 language model.

Architectural Choices and Model Modifications

Building a multimodal model from scratch demands massive compute power. A smarter approach is combining pre-trained components using Hugging Face's VisionEncoderDecoderModel architecture. We selected google/vit-base-patch16-224-in21k as the encoder for extracting visual features. The decoder is the Russian-language model ai-forever/rugpt3small_based_on_gpt2.

Key challenge: ruGPT-3 doesn't natively support cross-attention to receive data from the encoder. The fix involves tweaking the model config before initialization:

Google AdInline article slot
from transformers import AutoConfig, AutoModelForCausalLM, AutoModel

encoder = AutoModel.from_pretrained("google/vit-base-patch16-224-in21k")

decoder_config = AutoConfig.from_pretrained("ai-forever/rugpt3small_based_on_gpt2")
decoder_config.is_decoder = True
decoder_config.add_cross_attention = True

decoder = AutoModelForCausalLM.from_pretrained("ai-forever/rugpt3small_based_on_gpt2", config=decoder_config)

Once assembled, the model initializes new weights for cross-attention and projection layers, which get fine-tuned during training.

Data Preparation and Processing

No ready Russian datasets pair X-ray images with medical reports, so we got creative. We used the English Indiana University Chest X-Ray (IU X-Ray) dataset with ~7,500 images and reports. Reports were translated to Russian using the Helsinki-NLP/opus-mt-en-ru model right in the Kaggle environment.

Kaggle file handling has quirks. Standard image loaders often only see thumbnail previews. For proper mapping, we deeply scan the filesystem:

Google AdInline article slot
import os

file_map = {}
for root, dirs, files in os.walk('/kaggle/input'):
    for file in files:
        if file.lower().endswith(('.png', '.jpg', '.jpeg')):
            base_name = os.path.splitext(file)[0]
            file_map[base_name] = os.path.join(root, file)

This ensures reliable linking of CSV identifiers to actual image paths.

Training Process and Troubleshooting

Training ran on two NVIDIA T4 GPUs with mixed precision (fp16) and gradient accumulation (gradient_accumulation_steps=4) to fit a virtual batch size of 32 efficiently.

Common Issues and Fixes:

Google AdInline article slot
  • Checkpoint save crashes: Seq2SeqTrainer in some transformers versions fails on custom composite models. Fix: Disable intermediate saves (save_strategy="no") and manually handle generation params via GenerationConfig before final save.
  • Kaggle session timeouts: Prevent auto-shutdown during long training by running a simple browser console JavaScript snippet to simulate periodic clicks.

Results, Inference, and Deployment

After 15 epochs, the model generates coherent Russian medical reports, using professional terminology correctly. It identifies basic patterns like clear lungs, pneumothorax signs, and cardiomegaly. The small dataset leads to occasional hallucinations—like mentioning absent medical devices.

Basic Inference Example:

import torch
from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
from PIL import Image

model_id = "livadies/Russian-Radiologist-ruGPT-ViT"
model = VisionEncoderDecoderModel.from_pretrained(model_id)
feature_extractor = ViTImageProcessor.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)

image = Image.open("xray.jpg").convert("RGB")
pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values

generated_ids = model.generate(pixel_values, max_length=128, num_beams=4)
print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))

A Gradio demo runs on Hugging Face Spaces using CPU, with 10–15 second generation times. Full source code—including data pipelines and training notebooks—is in the public repo.

Key Takeaways

  • Hybrid Architecture: Successfully fused pre-trained ViT visual encoder and ruGPT-3 language model by adding cross-attention via config tweaks.
  • Data Workarounds: Overcame lack of Russian medical datasets by translating English ones and navigating Kaggle's file quirks.
  • Practical Fixes: Disabled buggy auto-saves in Seq2SeqTrainer and prevented session timeouts for reliable long training.
  • Proof of Concept: Small-dataset model validates Vision + Language architecture for structured Russian medical text generation.
  • Open and Accessible: Complete code and interactive demo are public for community study, testing, and extension.

— Editorial Team

Advertisement 728x90

Read Next