File size: 3,212 Bytes
857620c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | import os
from typing import Optional
from cog import BasePredictor, Input, Path
from llama_cpp import Llama
class Predictor(BasePredictor):
def setup(self) -> None:
"""Load the model into memory"""
model_path = "monad-mistral-7b.gguf"
# Download model if not present (Replicate will cache this)
if not os.path.exists(model_path):
print(f"Model not found at {model_path}")
# Replicate will handle model file placement
self.llm = Llama(
model_path=model_path,
n_ctx=4096, # Context window
n_threads=8, # CPU threads
n_gpu_layers=-1, # Use all GPU layers
verbose=False
)
# Default generation parameters
self.default_params = {
"max_tokens": 512,
"temperature": 0.7,
"top_p": 0.9,
"top_k": 40,
"repeat_penalty": 1.1
}
def predict(
self,
prompt: str = Input(
description="Input prompt",
default="What is Monad blockchain?"
),
system_prompt: str = Input(
description="System prompt to guide the model's behavior",
default="You are an expert on Monad blockchain technology. Provide accurate, helpful information about Monad's architecture, ecosystem, and capabilities."
),
max_tokens: int = Input(
description="Maximum number of tokens to generate",
default=512,
ge=1,
le=4096
),
temperature: float = Input(
description="Temperature for sampling",
default=0.7,
ge=0.1,
le=2.0
),
top_p: float = Input(
description="Top-p sampling parameter",
default=0.9,
ge=0.1,
le=1.0
),
top_k: int = Input(
description="Top-k sampling parameter",
default=40,
ge=1,
le=100
),
repeat_penalty: float = Input(
description="Penalty for repeated tokens",
default=1.1,
ge=1.0,
le=2.0
),
seed: int = Input(
description="Random seed for reproducibility",
default=-1
)
) -> str:
"""Run inference on the model"""
# Format prompt with Mistral template
if system_prompt:
formatted_prompt = f"[INST] {system_prompt}\n\n{prompt} [/INST]"
else:
formatted_prompt = f"[INST] {prompt} [/INST]"
# Set seed if provided
if seed > 0:
import random
import numpy as np
random.seed(seed)
np.random.seed(seed)
# Generate response
output = self.llm(
formatted_prompt,
max_tokens=max_tokens,
temperature=temperature,
top_p=top_p,
top_k=top_k,
repeat_penalty=repeat_penalty,
stop=["</s>", "[INST]", "[/INST]"],
echo=False
)
return output['choices'][0]['text'].strip() |