Ollama와 Qwen3 Reranker 모델을 사용한 문서 재정렬 - Go로

$ ./rnk ./example_query.txt ./example_docs

임베딩 모델 사용: dengcao/Qwen3-Embedding-4B:Q5_K_M
Ollama 기본 URL: http://localhost:11434
쿼리 파일 처리: ./example_query.txt, 대상 디렉토리: ./example_docs
쿼리: 인공지능이 무엇인지, 그리고 머신러닝이 어떻게 작동하는지?
7개 문서 발견
쿼리 임베딩 추출...
문서 처리...

=== 유사도에 따른 순위 ===
1. example_docs/ai_introduction.txt (점수: 0.451)
2. example_docs/machine_learning.md (점수: 0.388)
3. example_docs/qwen3-reranking-models.md (점수: 0.354)
4. example_docs/ollama-parallelism.md (점수: 0.338)
5. example_docs/ollama-reranking-models.md (점수: 0.318)
6. example_docs/programming_basics.txt (점수: 0.296)
7. example_docs/setup.log (점수: 0.282)

7개 문서 처리 완료, 소요 시간: 2.023초 (평균: 0.289초/문서)
재정렬기 모델을 사용하여 문서 재정렬...
dengcao/Qwen3-Reranker-4B:Q5_K_M을 사용하여 cross-encoder 접근 방식으로 재정렬 수행

=== 재정렬기로 순위 매기기 ===
1. example_docs/ai_introduction.txt (점수: 0.343)
2. example_docs/machine_learning.md (점수: 0.340)
3. example_docs/programming_basics.txt (점수: 0.320)
4. example_docs/setup.log (점수: 0.313)
5. example_docs/ollama-parallelism.md (점수: 0.313)
6. example_docs/qwen3-reranking-models.md (점수: 0.312)
7. example_docs/ollama-reranking-models.md (점수: 0.306)

7개 문서 처리 완료, 소요 시간: 1.984초 (평균: 0.283초/문서)

startTime = time.Now()
// 재정렬기 모델을 사용하여 재정렬
fmt.Println("재정렬기 모델을 사용하여 문서 재정렬 중...")

// rerankingModel := "dengcao/Qwen3-Reranker-0.6B:F16"
rerankingModel := "dengcao/Qwen3-Reranker-4B:Q5_K_M"
rerankedDocs, err := rerankDocuments(validDocs, query, rerankingModel, ollamaBaseURL)
if err != nil {
	log.Fatalf("문서 재정렬 중 오류 발생: %v", err)
}

fmt.Println("\n=== 재정렬기로 순위 매기기 ===")
for i, doc := range rerankedDocs {
	fmt.Printf("%d. %s (점수: %.3f)\n", i+1, doc.Path, doc.Score)
}

totalTime = time.Since(startTime)
avgTimePerDoc = totalTime / time.Duration(len(rerankedDocs))

fmt.Printf("\n%d개 문서 처리 완료, 소요 시간: %.3fs (평균: %.3fs/문서)\n",
	len(rerankedDocs), totalTime.Seconds(), avgTimePerDoc.Seconds())

func rerankDocuments(validDocs []Document, query, rerankingModel, ollamaBaseURL string) ([]Document, error) {
	// 표준 Ollama에는 직접적인 재정렬 API가 없기 때문에, 쿼리-문서 쌍의 임베딩을 생성하고 이를 점수화하여
	// 재정렬을 구현해야 합니다

	fmt.Println("cross-encoder 접근 방식으로", rerankingModel, "을 사용하여 재정렬 수행 중")

	rerankedDocs := make([]Document, len(validDocs))
	copy(rerankedDocs, validDocs)

	for i, doc := range validDocs {
		// 쿼리와 문서를 결합하여 재정렬을 위한 프롬프트 생성
		rerankPrompt := fmt.Sprintf("Query: %s\n\nDocument: %s\n\nRelevance:", query, doc.Content)

		// 결합된 프롬프트의 임베딩 가져오기
		embedding, err := getEmbedding(rerankPrompt, rerankingModel, ollamaBaseURL)
		if err != nil {
			fmt.Printf("경고: 문서 %d에 대한 재정렬 임베딩 가져오기 실패: %v\n", i, err)
			// 중립 점수로 대체
			rerankedDocs[i].Score = 0.5
			continue
		}

		// 임베딩의 크기를 사용하여 관련성 점수 계산
		// (이 방법은 간단한 접근 방식이며, 실제로는 훈련된 재정렬기 사용이 필요함)
		score := calculateRelevanceScore(embedding)
		rerankedDocs[i].Score = score
		// fmt.Printf("문서 %d 재정렬 완료, 점수: %.4f\n", i, score)
	}

	// 재정렬 점수(내림차순)에 따라 문서 정렬
	sort.Slice(rerankedDocs, func(i, j int) bool {
		return rerankedDocs[i].Score > rerankedDocs[j].Score
	})

	return rerankedDocs, nil
}

func calculateRelevanceScore(embedding []float64) float64 {
	// 임베딩 크기와 양수 값에 기반한 간단한 점수 계산
	var sumPositive, sumTotal float64
	for _, val := range embedding {
		sumTotal += val * val
		if val > 0 {
			sumPositive += val
		}
	}

	if sumTotal == 0 {
		return 0
	}

	// 크기와 양수 편향을 정규화하여 결합
	magnitude := math.Sqrt(sumTotal) / float64(len(embedding))
	positiveRatio := sumPositive / float64(len(embedding))

	return (magnitude + positiveRatio) / 2
}

import (
	"math"
)

go build -o rnk

./rnk ./example_query.txt ./example_docs