107 lines
3.2 KiB
Java
107 lines
3.2 KiB
Java
package org.tzd.lm;
|
||
|
||
import com.axis.innovators.box.tools.FolderCreator;
|
||
import com.axis.innovators.box.tools.LibraryLoad;
|
||
|
||
/**
|
||
* LM推理类
|
||
* @author tzdwindows 7
|
||
*/
|
||
public class LM {
|
||
public static boolean CUDA = false;
|
||
public final static String DEEP_SEEK = FolderCreator.getModelFolder() + "//DeepSeek-R1-Distill-Qwen-1.5B-Q8_0.gguf";
|
||
|
||
static {
|
||
if (!CUDA) {
|
||
LibraryLoad.loadLibrary("cpu/ggml-base");
|
||
LibraryLoad.loadLibrary("cpu/ggml-cpu");
|
||
LibraryLoad.loadLibrary("cpu/ggml");
|
||
LibraryLoad.loadLibrary("cpu/llama");
|
||
} else {
|
||
LibraryLoad.loadLibrary("cuda/ggml-base");
|
||
LibraryLoad.loadLibrary("cuda/ggml-cpu");
|
||
LibraryLoad.loadLibrary("cuda/ggml-rpc");
|
||
// cuda版本 cuda-cu12.4-x64（确保你有）
|
||
LibraryLoad.loadLibrary("cuda/ggml-cuda");
|
||
LibraryLoad.loadLibrary("cuda/ggml");
|
||
LibraryLoad.loadLibrary("cuda/llama");
|
||
}
|
||
LibraryLoad.loadLibrary("LM");
|
||
}
|
||
/**
|
||
* 加载模型
|
||
* @param pathModel 模型路径
|
||
* @return 模型句柄
|
||
*/
|
||
public static native long llamaLoadModelFromFile(String pathModel);
|
||
|
||
/**
|
||
* 释放模型资源
|
||
* @param modelHandle 模型句柄
|
||
*/
|
||
public static native void llamaFreeModel(long modelHandle);
|
||
|
||
/**
|
||
* 上下文创建
|
||
* @param modelHandle 上下文句柄
|
||
* @return 上下文句柄
|
||
*/
|
||
public static native long createContext(long modelHandle);
|
||
|
||
/**
|
||
* 释放上下文资源
|
||
* @param ctxHandle 上下文句柄
|
||
*/
|
||
public static native void llamaFreeContext(long ctxHandle);
|
||
|
||
/**
|
||
* 推理模型
|
||
* @param modelHandle 模型句柄
|
||
* @param ctxHandle 模型上下文句柄
|
||
* @param temperature 温度
|
||
* @param prompt 问题
|
||
* @param messageCallback 回调接口
|
||
* @return 最终内容
|
||
*/
|
||
public static native String inference(long modelHandle ,
|
||
long ctxHandle,
|
||
float temperature,
|
||
String prompt,
|
||
MessageCallback messageCallback);
|
||
|
||
/**
|
||
* 回调接口
|
||
*/
|
||
public interface MessageCallback {
|
||
/**
|
||
* 接口回调
|
||
* @param message 消息
|
||
*/
|
||
void onMessage(String message);
|
||
}
|
||
|
||
public static void main(String[] args) {
|
||
// 加载模型
|
||
long modelHandle = llamaLoadModelFromFile(DEEP_SEEK);
|
||
// 创建新的上下文
|
||
long ctxHandle = createContext(modelHandle);
|
||
inference(modelHandle, ctxHandle, 0.2f, "写一个ai", new MessageCallback() {
|
||
@Override
|
||
public void onMessage(String message) {
|
||
// 回调输出
|
||
System.out.print(message);
|
||
}
|
||
});
|
||
// 推理模型
|
||
inference(modelHandle, ctxHandle, 0.2f, "谢谢你", new MessageCallback() {
|
||
@Override
|
||
public void onMessage(String message) {
|
||
// 回调输出
|
||
System.out.print(message);
|
||
}
|
||
});
|
||
// 清理上下文
|
||
llamaFreeContext(ctxHandle);
|
||
}
|
||
}
|