Add example to run Hugging Face models using Ollama (#8771)

Author: ilopezluna

examples/ollama-hugging-face/build.gradle

@@ -0,0 +1,19 @@
+plugins {
+    id 'java'
+}
+
+repositories {
+    mavenCentral()
+}
+
+dependencies {
+    testImplementation 'org.testcontainers:ollama'
+    testImplementation 'org.assertj:assertj-core:3.25.3'
+    testImplementation 'ch.qos.logback:logback-classic:1.3.14'
+    testImplementation 'org.junit.jupiter:junit-jupiter:5.10.2'
+    testImplementation 'io.rest-assured:rest-assured:5.4.0'
+}
+
+test {
+    useJUnitPlatform()
+}

examples/ollama-hugging-face/src/test/java/com/example/ollamahf/OllamaHuggingFaceContainer.java

@@ -0,0 +1,70 @@
+package com.example.ollamahf;
+
+import com.github.dockerjava.api.command.InspectContainerResponse;
+import org.testcontainers.containers.ContainerLaunchException;
+import org.testcontainers.ollama.OllamaContainer;
+import org.testcontainers.utility.DockerImageName;
+
+import java.io.IOException;
+
+public class OllamaHuggingFaceContainer extends OllamaContainer {
+
+    private final HuggingFaceModel huggingFaceModel;
+
+    public OllamaHuggingFaceContainer(HuggingFaceModel model) {
+        super(DockerImageName.parse("ollama/ollama:0.1.47"));
+        this.huggingFaceModel = model;
+    }
+
+    @Override
+    protected void containerIsStarted(InspectContainerResponse containerInfo, boolean reused) {
+        super.containerIsStarted(containerInfo, reused);
+        if (reused || huggingFaceModel == null) {
+            return;
+        }
+
+        try {
+            executeCommand("apt-get", "update");
+            executeCommand("apt-get", "upgrade", "-y");
+            executeCommand("apt-get", "install", "-y", "python3-pip");
+            executeCommand("pip", "install", "huggingface-hub");
+            executeCommand(
+                "huggingface-cli",
+                "download",
+                huggingFaceModel.repository,
+                huggingFaceModel.model,
+                "--local-dir",
+                "."
+            );
+            executeCommand("sh", "-c", String.format("echo '%s' > Modelfile", huggingFaceModel.modelfileContent));
+            executeCommand("ollama", "create", huggingFaceModel.model, "-f", "Modelfile");
+            executeCommand("rm", huggingFaceModel.model);
+        } catch (IOException | InterruptedException e) {
+            throw new ContainerLaunchException(e.getMessage());
+        }
+    }
+
+    private void executeCommand(String... command) throws ContainerLaunchException, IOException, InterruptedException {
+        ExecResult execResult = execInContainer(command);
+        if (execResult.getExitCode() > 0) {
+            throw new ContainerLaunchException(
+                "Failed to execute " + String.join(" ", command) + ": " + execResult.getStderr()
+            );
+        }
+    }
+
+    public static class HuggingFaceModel {
+
+        public final String repository;
+
+        public final String model;
+
+        public String modelfileContent;
+
+        public HuggingFaceModel(String repository, String model) {
+            this.repository = repository;
+            this.model = model;
+            this.modelfileContent = "FROM " + model;
+        }
+    }
+}

examples/ollama-hugging-face/src/test/java/com/example/ollamahf/OllamaHuggingFaceTest.java

@@ -0,0 +1,74 @@
+package com.example.ollamahf;
+
+import io.restassured.http.Header;
+import org.junit.jupiter.api.Test;
+import org.testcontainers.ollama.OllamaContainer;
+import org.testcontainers.utility.DockerImageName;
+
+import java.util.List;
+
+import static io.restassured.RestAssured.given;
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class OllamaHuggingFaceTest {
+
+    @Test
+    public void embeddingModelWithHuggingFace() {
+        String repository = "CompendiumLabs/bge-small-en-v1.5-gguf";
+        String model = "bge-small-en-v1.5-q4_k_m.gguf";
+        String imageName = "embedding-model-from-hugging-face";
+        OllamaContainer ollama = new OllamaContainer(
+            DockerImageName.parse(imageName).asCompatibleSubstituteFor("ollama/ollama:0.1.47")
+        );
+        boolean imageExists = ollama
+            .getDockerClient()
+            .listImagesCmd()
+            .exec()
+            .stream()
+            .anyMatch(image -> image.getRepoTags()[0].equals(imageName + ":latest"));
+        if (!imageExists) {
+            createImage(imageName, repository, model);
+        }
+        ollama.start();
+
+        String modelName = given()
+            .baseUri(ollama.getEndpoint())
+            .get("/api/tags")
+            .jsonPath()
+            .getString("models[0].name");
+        assertThat(modelName).contains(model + ":latest");
+
+        List<Float> embedding = given()
+            .baseUri(ollama.getEndpoint())
+            .header(new Header("Content-Type", "application/json"))
+            .body(new EmbeddingRequest(model + ":latest", "Hello from Testcontainers!"))
+            .post("/api/embeddings")
+            .jsonPath()
+            .getList("embedding");
+
+        assertThat(embedding).isNotNull();
+        assertThat(embedding.isEmpty()).isFalse();
+    }
+
+    private static void createImage(String imageName, String repository, String model) {
+        OllamaHuggingFaceContainer.HuggingFaceModel hfModel = new OllamaHuggingFaceContainer.HuggingFaceModel(
+            repository,
+            model
+        );
+        OllamaHuggingFaceContainer huggingFaceContainer = new OllamaHuggingFaceContainer(hfModel);
+        huggingFaceContainer.start();
+        huggingFaceContainer.commitToImage(imageName);
+    }
+
+    public static class EmbeddingRequest {
+
+        public final String model;
+
+        public final String prompt;
+
+        public EmbeddingRequest(String model, String prompt) {
+            this.model = model;
+            this.prompt = prompt;
+        }
+    }
+}

examples/ollama-hugging-face/src/test/resources/logback-test.xml

@@ -0,0 +1,16 @@
+<configuration>
+
+    <appender name="STDOUT" class="ch.qos.logback.core.ConsoleAppender">
+        <!-- encoders are assigned the type
+             ch.qos.logback.classic.encoder.PatternLayoutEncoder by default -->
+        <encoder>
+            <pattern>%d{HH:mm:ss.SSS} %-5level %logger - %msg%n</pattern>
+        </encoder>
+    </appender>
+
+    <root level="INFO">
+        <appender-ref ref="STDOUT"/>
+    </root>
+
+    <logger name="org.testcontainers" level="INFO"/>
+</configuration>

examples/settings.gradle

@@ -35,6 +35,7 @@ include 'zookeeper'
 include 'hazelcast'
 include 'nats'
 include 'sftp'
+include 'ollama-hugging-face'
 
 ext.isCI = System.getenv("CI") != null