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:
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:
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:
- 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
No comments yet.