| from typing import Dict, List, Any | |
| from llama_cpp import Llama | |
| import torch | |
| MAX_TOKENS=8192 | |
| GPU_LAYERS=99 | |
| class EndpointHandler(): | |
| def __init__(self, data): | |
| n_gpu_layers = GPU_LAYERS | |
| if not torch.cuda.is_available(): | |
| n_gpu_layers = 0 | |
| self.model = Llama.from_pretrained("lmstudio-ai/gemma-2b-it-GGUF", filename="gemma-2b-it-q4_k_m.gguf", n_ctx=8192, cache_dir="./", n_gpu_layers=n_gpu_layers) | |
| def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: | |
| inputs = data.pop("inputs", "") | |
| temperature = data.pop("temperature", None) | |
| if not temperature: | |
| temperature = data.pop("temp", 0.33) | |
| if temperature > 3 or temperature < 0: | |
| return json.dumps({ | |
| "status": "error", | |
| "reason": "invalid temperature ( 0.01 - 1.00 )" | |
| }) | |
| top_p = data.pop("top-p", 0.85) | |
| if top_p > 3 or top_p < 0: | |
| return json.dumps({ | |
| "status": "error", | |
| "reason": "invalid top percentage ( 0.01 - 1.00 )" | |
| }) | |
| top_k = data.pop("top-k", 42) | |
| if top_k > 100 or top_k < 0: | |
| return json.dumps({ | |
| "status": "error", | |
| "reason": "invalid top k ( 1 - 99 )" | |
| }) | |
| system_prompt = data.pop("system-prompt", "You are Gemma. Assist user with whatever they require, in a safe and moral manner.") | |
| format = data.pop("format", "<startofturn>system\n{system_prompt} <endoftext>\n<startofturn>user\n{prompt} <endofturn>\n<startofturn>model") | |
| try: | |
| format = format.format(system_prompt = system_prompt, prompt = inputs) | |
| except Exception as e: | |
| return json.dumps({ | |
| "status": "error", | |
| "reason": "invalid format" | |
| }) | |
| res = self.model(format, temperature=temperature, top_p=top_p, top_k=42) | |
| return res |