removed std print and adding ranking test.

Vaijanath Rao · Vaijanath Rao · commit faa494e88682 · 2025-03-13T15:41:56.000-07:00
diff --git a/src/main/cpp/jllama.cpp b/src/main/cpp/jllama.cpp
@@ -765,8 +765,6 @@ JNIEXPORT jobject JNICALL Java_de_kherud_llama_LlamaModel_rerank(JNIEnv *env, jo
     
 		const auto out_res = result->to_json();
     	
-		std::cout << out_res.dump(4) << std::endl;
-    	
 		if (result->is_stop()) {
 			for (const int id_task : task_ids) {
 				ctx_server->queue_results.remove_waiting_task_id(id_task);
diff --git a/src/main/java/de/kherud/llama/LlamaIterator.java b/src/main/java/de/kherud/llama/LlamaIterator.java
@@ -35,6 +35,9 @@ public LlamaOutput next() {
         }
         LlamaOutput output = model.receiveCompletion(taskId);
         hasNext = !output.stop;
+        if (output.stop) {
+        	model.releaseTask(taskId);
+        }
         return output;
     }
 
diff --git a/src/main/java/de/kherud/llama/LlamaModel.java b/src/main/java/de/kherud/llama/LlamaModel.java
@@ -5,7 +5,9 @@
 
 import java.lang.annotation.Native;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.function.BiConsumer;
 
 /**
@@ -131,13 +133,34 @@ public void close() {
 
 	private native void delete();
 	
-	private native void releaseTask(int taskId);
+	native void releaseTask(int taskId);
 
 	private static native byte[] jsonSchemaToGrammarBytes(String schema);
 	
 	public static String jsonSchemaToGrammar(String schema) {
 		return new String(jsonSchemaToGrammarBytes(schema), StandardCharsets.UTF_8);
 	}
 	
+	public List<Pair<String, Float>> rerank(boolean reRank, String query, String ... documents) {
+		LlamaOutput output = rerank(query, documents);
+		
+		Map<String, Float> scoredDocumentMap = output.probabilities;
+		
+		List<Pair<String, Float>> rankedDocuments = new ArrayList<>();
+		
+		if (reRank) {
+            // Sort in descending order based on Float values
+            scoredDocumentMap.entrySet()
+                    .stream()
+                    .sorted((a, b) -> Float.compare(b.getValue(), a.getValue())) // Descending order
+                    .forEach(entry -> rankedDocuments.add(new Pair<>(entry.getKey(), entry.getValue())));
+        } else {
+            // Copy without sorting
+            scoredDocumentMap.forEach((key, value) -> rankedDocuments.add(new Pair<>(key, value)));
+        }
+		
+		return rankedDocuments;
+	}
+	
 	public native LlamaOutput rerank(String query, String... documents);
 }
diff --git a/src/main/java/de/kherud/llama/Pair.java b/src/main/java/de/kherud/llama/Pair.java
@@ -0,0 +1,48 @@
+package de.kherud.llama;
+
+import java.util.Objects;
+
+public class Pair<K, V> {
+
+	private final K key;
+	private final V value;
+	
+	public Pair(K key, V value) {
+		this.key = key;
+		this.value = value;
+	}
+	
+	public K getKey() {
+		return key;
+	}
+	
+	public V getValue() {
+		return value;
+	}
+
+	@Override
+	public int hashCode() {
+		return Objects.hash(key, value);
+	}
+
+	@Override
+	public boolean equals(Object obj) {
+		if (this == obj)
+			return true;
+		if (obj == null)
+			return false;
+		if (getClass() != obj.getClass())
+			return false;
+		Pair other = (Pair) obj;
+		return Objects.equals(key, other.key) && Objects.equals(value, other.value);
+	}
+
+	@Override
+	public String toString() {
+		return "Pair [key=" + key + ", value=" + value + "]";
+	}
+	
+	
+	
+	
+}
diff --git a/src/test/java/de/kherud/llama/RerankingModelTest.java b/src/test/java/de/kherud/llama/RerankingModelTest.java
@@ -1,5 +1,6 @@
 package de.kherud.llama;
 
+import java.util.List;
 import java.util.Map;
 
 import org.junit.AfterClass;
@@ -10,6 +11,13 @@
 public class RerankingModelTest {
 
 	private static LlamaModel model;
+	
+	String query = "Machine learning is";
+	String[] TEST_DOCUMENTS = new String[] {
+			"A machine is a physical system that uses power to apply forces and control movement to perform an action. The term is commonly applied to artificial devices, such as those employing engines or motors, but also to natural biological macromolecules, such as molecular machines.",
+			"Learning is the process of acquiring new understanding, knowledge, behaviors, skills, values, attitudes, and preferences. The ability to learn is possessed by humans, non-human animals, and some machines; there is also evidence for some kind of learning in certain plants.",
+			"Machine learning is a field of study in artificial intelligence concerned with the development and study of statistical algorithms that can learn from data and generalize to unseen data, and thus perform tasks without explicit instructions.",
+			"Paris, capitale de la France, est une grande ville européenne et un centre mondial de l'art, de la mode, de la gastronomie et de la culture. Son paysage urbain du XIXe siècle est traversé par de larges boulevards et la Seine." };
 
 	@BeforeClass
 	public static void setup() {
@@ -28,12 +36,7 @@ public static void tearDown() {
 	@Test
 	public void testReRanking() {
 
-		String query = "Machine learning is";
-		String[] TEST_DOCUMENTS = new String[] {
-				"A machine is a physical system that uses power to apply forces and control movement to perform an action. The term is commonly applied to artificial devices, such as those employing engines or motors, but also to natural biological macromolecules, such as molecular machines.",
-				"Learning is the process of acquiring new understanding, knowledge, behaviors, skills, values, attitudes, and preferences. The ability to learn is possessed by humans, non-human animals, and some machines; there is also evidence for some kind of learning in certain plants.",
-				"Machine learning is a field of study in artificial intelligence concerned with the development and study of statistical algorithms that can learn from data and generalize to unseen data, and thus perform tasks without explicit instructions.",
-				"Paris, capitale de la France, est une grande ville européenne et un centre mondial de l'art, de la mode, de la gastronomie et de la culture. Son paysage urbain du XIXe siècle est traversé par de larges boulevards et la Seine." };
+		
 		LlamaOutput llamaOutput = model.rerank(query, TEST_DOCUMENTS[0], TEST_DOCUMENTS[1], TEST_DOCUMENTS[2],
 				TEST_DOCUMENTS[3]);
 
@@ -64,5 +67,17 @@ public void testReRanking() {
 
 		
 	}
-
+	
+	@Test
+	public void testSortedReRanking() {
+		List<Pair<String, Float>> rankedDocuments = model.rerank(true, query, TEST_DOCUMENTS);
+		Assert.assertEquals(rankedDocuments.size(), TEST_DOCUMENTS.length);
+		
+		// Check the ranking order: each score should be >= the next one
+	    for (int i = 0; i < rankedDocuments.size() - 1; i++) {
+	        float currentScore = rankedDocuments.get(i).getValue();
+	        float nextScore = rankedDocuments.get(i + 1).getValue();
+	        Assert.assertTrue("Ranking order incorrect at index " + i, currentScore >= nextScore);
+	    }
+	}
 }

Original file line number	Diff line number	Diff line change
`@@ -35,6 +35,9 @@ public LlamaOutput next() {`
`35`	`35`	`}`
`36`	`36`	`LlamaOutput output = model.receiveCompletion(taskId);`
`37`	`37`	`hasNext = !output.stop;`
	`38`	`+ if (output.stop) {`
	`39`	`+ model.releaseTask(taskId);`
	`40`	`+ }`
`38`	`41`	`return output;`
`39`	`42`	`}`
`40`	`43`