Reordenar documentos con Ollama y el modelo Qwen3 Reranker - en Go

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

Usando modelo de embedding: dengcao/Qwen3-Embedding-4B:Q5_K_M
URL base de Ollama: http://localhost:11434
Procesando archivo de consulta: ./example_query.txt, directorio objetivo: ./example_docs
Consulta: ¿Qué es la inteligencia artificial y cómo funciona el aprendizaje automático?
Se encontraron 7 documentos
Extrayendo embedding de consulta...
Procesando documentos...

=== CLASIFICACIÓN POR SIMILITUD ===
1. example_docs/ai_introduction.txt (Puntuación: 0,451)
2. example_docs/machine_learning.md (Puntuación: 0,388)
3. example_docs/qwen3-reranking-models.md (Puntuación: 0,354)
4. example_docs/ollama-parallelism.md (Puntuación: 0,338)
5. example_docs/ollama-reranking-models.md (Puntuación: 0,318)
6. example_docs/programming_basics.txt (Puntuación: 0,296)
7. example_docs/setup.log (Puntuación: 0,282)

Procesados 7 documentos en 2,023s (promedio: 0,289s por documento)
Clasificando documentos con modelo reranker...
Implementando clasificación usando enfoque cross-encoder con dengcao/Qwen3-Reranker-4B:Q5_K_M

=== CLASIFICACIÓN CON RERANKER ===
1. example_docs/ai_introduction.txt (Puntuación: 0,343)
2. example_docs/machine_learning.md (Puntuación: 0,340)
3. example_docs/programming_basics.txt (Puntuación: 0,320)
4. example_docs/setup.log (Puntuación: 0,313)
5. example_docs/ollama-parallelism.md (Puntuación: 0,313)
6. example_docs/qwen3-reranking-models.md (Puntuación: 0,312)
7. example_docs/ollama-reranking-models.md (Puntuación: 0,306)

Procesados 7 documentos en 1,984s (promedio: 0,283s por documento)

  startTime = time.Now()
	// rerank usando modelo reranker
	fmt.Println("Clasificando documentos con modelo reranker...")

	// 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("Error al clasificar documentos: %v", err)
	}

	fmt.Println("\n=== CLASIFICACIÓN CON RERANKER ===")
	for i, doc := range rerankedDocs {
		fmt.Printf("%d. %s (Puntuación: %.3f)\n", i+1, doc.Path, doc.Score)
	}

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

	fmt.Printf("\nProcesados %d documentos en %.3fs (promedio: %.3fs por documento)\n",
		len(rerankedDocs), totalTime.Seconds(), avgTimePerDoc.Seconds())

func rerankDocuments(validDocs []Document, query, rerankingModel, ollamaBaseURL string) ([]Document, error) {
	// Dado que el Ollama estándar no tiene una API de rerank directa, implementaremos
	// el reranking generando embeddings para pares de consulta y documento y calificándolos

	fmt.Println("Implementando clasificación usando enfoque cross-encoder con", rerankingModel)

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

	for i, doc := range validDocs {
		// Crear un prompt para reranking combinando consulta y documento
		rerankPrompt := fmt.Sprintf("Consulta: %s\n\nDocumento: %s\n\nRelevancia:", query, doc.Content)

		// Obtener embedding para el prompt combinado
		embedding, err := getEmbedding(rerankPrompt, rerankingModel, ollamaBaseURL)
		if err != nil {
			fmt.Printf("Advertencia: No se pudo obtener el embedding de clasificación para el documento %d: %v\n", i, err)
			// Retroceso a una puntuación neutral
			rerankedDocs[i].Score = 0.5
			continue
		}

		// Usar la magnitud del embedding como puntuación de relevancia
		// (Este es un enfoque simplificado - en la práctica, usarías un reranker entrenado)
		score := calculateRelevanceScore(embedding)
		rerankedDocs[i].Score = score
		// fmt.Printf("Documento %d clasificado con puntuación: %.4f\n", i, score)
	}

	// Clasificar documentos por puntuación de reranking (de mayor a menor)
	sort.Slice(rerankedDocs, func(i, j int) bool {
		return rerankedDocs[i].Score > rerankedDocs[j].Score
	})

	return rerankedDocs, nil
}

func calculateRelevanceScore(embedding []float64) float64 {
	// Puntuación simple basada en magnitud del embedding y valores positivos
	var sumPositive, sumTotal float64
	for _, val := range embedding {
		sumTotal += val * val
		if val > 0 {
			sumPositive += val
		}
	}

	if sumTotal == 0 {
		return 0
	}

	// Normalizar y combinar magnitud con sesgo positivo
	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