Kwaipilot OASIS-1.3B：代码检索新范式，语义理解突破行业基准-优快云博客

Kwaipilot OASIS-1.3B：代码检索新范式，语义理解突破行业基准

【免费下载链接】OASIS-code-1.3B 项目地址: https://ai.gitcode.com/hf_mirrors/Kwaipilot/OASIS-code-1.3B

[![Kwaipilot OASIS模型标志图](https://raw.githubusercontent.com/Anditty/OASIS/refs/heads/main/Group.svg)](https://link.gitcode.com/i/a2f9c7557c1ee3fcca6147d04e359226)

如上图所示，该标志图直观呈现了Kwaipilot OASIS模型的品牌视觉形象。这一设计元素不仅强化了模型的技术识别度，更为开发者提供了直观的品牌联想，帮助用户快速建立对OASIS系列模型的认知。

模型架构解析

项目代号：OASIS（Order-Augmented Strategy for Improved code Search）

技术特性

作为Kwaipilot团队研发的新一代代码嵌入模型，OASIS通过三大核心创新突破传统技术瓶颈：首创仓库级程序结构分析机制，实现代码上下文深度理解；自研OASIS-instruct数据合成引擎，构建高质量训练语料；设计多模态融合损失函数，显著提升语义匹配精度，全面刷新代码检索领域性能指标。

应用场景

该模型专为智能代码检索系统优化，特别适用于以下场景：开源项目代码片段精准匹配、IDE集成式代码推荐、大规模代码库语义化管理。其核心优势在于跨语言、跨风格的代码语义理解能力，有效降低开发者搜索与复用代码的时间成本。

训练体系

OASIS模型训练数据源自千万级开源仓库的结构化分析，覆盖20+编程语言及主流编码规范。通过自监督学习与领域知识蒸馏相结合的训练策略，模型在保持1.3B参数量级的同时，实现了精度与效率的最优平衡。

技术演进路线

Kwaipilot实验室未来研发规划包括：

✅ 已完成：高级别模型开源（最新版OASIS-code-1.5B已发布）
✅ 已完成：技术文档公开（arxiv论文全面阐述模型原理）
进行中：通用自然语言处理模型研发
规划中：多模态代码理解系统构建
探索中：垂直领域代码生成解决方案

性能评估报告

模型名称	参数量	CoSQA	AdvTest	CSN-Py	CSN-Ja	CSN-JS	CSN-PHP	CSN-Go	CSN-Ruby	平均得分
Openai-Embedding-Ada-002	未公开	0.4423	0.3808	0.6802	0.7149	0.6750	0.6062	0.8563	0.7472	0.6378
jina-embeddings-v2-base-code	161M	0.6837	0.3850	0.6634	0.6803	0.6304	0.5701	0.8595	0.7095	0.6477
CodeSage-large	1.3B	0.4753	0.5267	0.7077	0.7021	0.6950	0.6133	0.8371	0.7192	0.6595
CodeFuse-CGE-Small	3.8B	0.5619	0.4639	0.6958	0.6863	0.6564	0.6133	0.8637	0.7341	0.6594
OASIS-1.3B	1.3B	0.5532	0.4861	0.7110	0.7199	0.6727	0.6217	0.8732	0.7333	0.6713

注：表格中加粗数据表示该指标当前最优值，平均得分为各数据集指标算术平均值

快速上手指南

基础调用方式

pip install -U torch transformers

兼容性提示：使用torch_dtype=torch.bfloat16加载模型时，请避免torch 2.5.0版本。推荐使用PyTorch 2.4.1及以下版本，或升级至2.5.1及以上版本以获得最佳稳定性。

import torch
import torch.nn.functional as F
from torch import Tensor
from transformers import AutoModel, AutoTokenizer

def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
    """句向量池化函数：提取最后一层隐藏状态的有效token特征"""
    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
    if left_padding:
        return last_hidden_states[:, -1]
    else:
        sequence_lengths = attention_mask.sum(dim=1) - 1
        batch_size = last_hidden_states.shape[0]
        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]

def get_query_prompt(query: str):
    """构建查询提示模板"""
    query_description = '根据代码搜索请求，检索最相关的代码片段'
    return f'Instruct: {query_description}\nQuery: {query}'

# 示例数据
search_query = "Python快速排序实现方法"
code_sample1 = """def bubble_sort(arr):
    n = len(arr)
    for i in range(n):
        swapped = False
        for j in range(1, n - i):
            if arr[j - 1] > arr[j]:
                arr[j - 1], arr[j] = arr[j], arr[j - 1]
                swapped = True
        if not swapped:
            break
    return arr"""

code_sample2 = """def quick_sort(arr):
    if len(arr) <= 1:
        return arr
    else:
        pivot = arr[0]
        less = [x for x in arr[1:] if x <= pivot]
        greater = [x for x in arr[1:] if x > pivot]
        return quick_sort(less) + [pivot] + quick_sort(greater)"""

# 模型加载
model = AutoModel.from_pretrained("Kwaipilot/OASIS-code-1.3B", output_hidden_states=True)
tokenizer = AutoTokenizer.from_pretrained("Kwaipilot/OASIS-code-1.3B")

# 文本编码与模型推理
inputs = tokenizer(
    [get_query_prompt(search_query), code_sample1, code_sample2],
    max_length=8192, 
    padding=True, 
    truncation=True, 
    return_tensors='pt'
)
outputs = model(**inputs)

# 特征提取与相似度计算
embeddings = last_token_pool(outputs.hidden_states[-1], inputs['attention_mask'])
print(f"嵌入向量维度: {embeddings.shape}")  # torch.Size([3, 2048])

# 余弦相似度计算
normalized_embeddings = F.normalize(embeddings, dim=1, p=2)
similarity_matrix = normalized_embeddings @ normalized_embeddings.T
print(f"查询与代码相似度: {similarity_matrix[0, 1:]}")  # tensor([0.6495, 0.8036])

Sentence Transformers集成方案

首先安装依赖库：

pip install -U sentence-transformers

通过Sentence Transformers API快速调用：

from sentence_transformers import SentenceTransformer

# 模型加载（支持bfloat16精度优化）
model = SentenceTransformer(
    "Kwaipilot/OASIS-code-1.3B",
    model_kwargs={"torch_dtype": torch.bfloat16}  # 可选精度设置
)

# 查询与代码示例
query = "Python快速排序实现方法"
code_snippet1 = """def bubble_sort(arr):
    n = len(arr)
    for i in range(n):
        swapped = False
        for j in range(1, n - i):
            if arr[j - 1] > arr[j]:
                arr[j - 1], arr[j] = arr[j], arr[j - 1]
                swapped = True
        if not swapped:
            break
    return arr"""

code_snippet2 = """def quick_sort(arr):
    if len(arr) <= 1:
        return arr
    else:
        pivot = arr[0]
        less = [x for x in arr[1:] if x <= pivot]
        greater = [x for x in arr[1:] if x > pivot]
        return quick_sort(less) + [pivot] + quick_sort(greater)"""

# 生成嵌入向量
query_embedding = model.encode([query], prompt_name="query")
code_embeddings = model.encode([code_snippet1, code_snippet2])

print(f"代码嵌入维度: {code_embeddings.shape}")  # (2, 2048)

# 相似度计算
print(f"查询-代码1相似度: {model.similarity(query_embedding[0], code_embeddings[0])}")  # tensor([[0.6495]])
print(f"查询-代码2相似度: {model.similarity(query_embedding[0], code_embeddings[1])}")  # tensor([[0.8036]])

学术引用格式

@misc{kwaipilotoasis,
  title = {Optimized Augmentation Strategy for Improved code Search},
  author = {Kwaipilot team},
  year = {2024},
}

OASIS系列模型通过创新的技术架构与系统的工程实践，正在重新定义代码理解的技术标准。随着1.5B版本的发布与NLP模型的研发推进，Kwaipilot团队将持续推动程序理解AI技术的边界，为开发者打造更智能、更高效的编码辅助工具。

【免费下载链接】OASIS-code-1.3B 项目地址: https://ai.gitcode.com/hf_mirrors/Kwaipilot/OASIS-code-1.3B

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

Kwaipilot OASIS-1.3B：代码检索新范式，语义理解突破行业基准