solbi12/ax4-mongodb-query-generator

pip install transformers torch

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 모델과 토크나이저 로드
model_name = "solbi12/ax4-mongodb-query-generator"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

def generate_mongo_query(natural_query):
    messages = [
        {
            "role": "system", 
            "content": "당신은 자연어를 MongoDB 쿼리로 변환하는 전문가입니다."
        },
        {
            "role": "user", 
            "content": natural_query
        }
    ]
    
    input_ids = tokenizer.apply_chat_template(
        messages, 
        add_generation_prompt=True, 
        return_tensors="pt"
    )
    
    with torch.no_grad():
        output = model.generate(
            input_ids,
            max_new_tokens=128,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id
        )
    
    response = tokenizer.decode(
        output[0][len(input_ids[0]):], 
        skip_special_tokens=True
    )
    return response

# 사용 예시
query = generate_mongo_query("가격이 5만원 이하인 상품들을 보여줘")
print(query)
# 출력: db.product.find({price: {$lte: 50000}})

queries = [
    "이번 달 매출 상위 브랜드는?",
    "고객 만족도가 높은 상품들",
    "재구매율이 높은 고객 세그먼트"
]

queries = [
    "지역별 주문 패턴 분석",
    "계절별 인기 카테고리",
    "마케팅 캠페인 효과 측정"
]

queries = [
    "오늘의 실시간 주문 현황",
    "재고 부족 상품 알림",
    "고객 서비스 우선순위"
]

LoRA Configuration:
  r: 16
  lora_alpha: 32
  lora_dropout: 0.1
  target_modules: ["q_proj", "v_proj", "k_proj", "o_proj"]

Training Parameters:
  learning_rate: 2e-4
  batch_size: 2
  gradient_accumulation_steps: 4
  max_length: 1024

@misc{solbi2024ax4mongodb,
  title={A.X-4.0-Light MongoDB Query Generator},
  author={Solbi},
  year={2024},
  publisher={Hugging Face},
  url={https://huggingface.co/solbi12/ax4-mongodb-query-generator}
}