Spaces:
Sleeping
Sleeping
File size: 1,347 Bytes
780954b |
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 |
import os
from typing import Optional
from transformers import pipeline
from .base import AbstractLLMModel
from .registry import register_llm_model
hf_token = os.getenv("HF_TOKEN")
@register_llm_model("meta-llama/Llama-")
class LlamaLLM(AbstractLLMModel):
def __init__(
self, model_id: str, device: str = "auto", cache_dir: str = "cache", **kwargs
):
super().__init__(model_id, device, cache_dir, **kwargs)
model_kwargs = kwargs.setdefault("model_kwargs", {})
model_kwargs["cache_dir"] = cache_dir
self.pipe = pipeline(
"text-generation",
model=model_id,
device_map=device,
return_full_text=False,
token=hf_token,
trust_remote_code=True,
**kwargs,
)
def generate(
self,
prompt: str,
system_prompt: Optional[
str
] = "You are a pirate chatbot who always responds in pirate speak!",
max_new_tokens: int = 256,
**kwargs
) -> str:
messages = []
if system_prompt:
messages.append({"role": "system", "content": system_prompt})
messages.append({"role": "user", "content": prompt})
outputs = self.pipe(messages, max_new_tokens=max_new_tokens, **kwargs)
return outputs[0]["generated_text"]
|