Enhancing Code Translation in Language Models with Few-Shot Learning via Retrieval-Augmented Generation
Manish Bhattarai, Javier E. Santos, Shawn M. Jones, Ayan Biswas, Boian Alexandrov, Daniel O’Malley · 2024
The advent of large language models (LLMs) has revolutionized the field of code translation, enabling automated translation between programming languages. Despite these advancements, the accuracy and reliability of these models often falter in complex translation tasks due to a lack of contextual understanding. This paper introduces a novel approach to enhance code translation through Few-Shot Learning augmented with retrieval-based techniques. By leveraging a repository of existing code translations, we dynamically retrieve the most relevant examples to guide the model in translating new code segments. Our method, based on Retrieval-Augmented Generation (RAG), significantly improves translation quality by providing contextual examples that the model can learn from in real-time. We chose RAG over traditional fine-tuning methods due to its ability to leverage existing codebases or a locally stored corpus of code, allowing it to dynamically adapt to diverse translation tasks without the need for extensive retraining. Extensive experiments on diverse datasets, using open LLM models such as Starcoder, Llama3-70B Instruct, CodeLlama-34B Instruct, Granite-34B Code Instruct, and Mixtral-8×22B, and commercial LLM models such as GPT-3.5 turbo, and GPT-4o demonstrate the superiority of our approach over traditional zero-shot, particularly in translating between Fortran and C++. We also explored different numbers of shots (examples provided to the model during inference) – specifically 1, 2, and 3 shots – and various embedding models for RAG, including Nomic-Embed, Starencoder, and CodeBERT, to evaluate the robustness and effectiveness of our approach.