File size: 2,513 Bytes
d31a6c6 |
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 |
from typing import List, Optional
from cog import BasePredictor, Input
from transformers import LLaMAForCausalLM, LLaMATokenizer
import torch
from train import PROMPT_DICT
PROMPT = PROMPT_DICT['prompt_no_input']
CACHE_DIR = 'alpaca_out'
class Predictor(BasePredictor):
def setup(self):
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
self.model = LLaMAForCausalLM.from_pretrained("alpaca_out", cache_dir=CACHE_DIR, local_files_only=True)
self.model = self.model
self.model.to(self.device)
self.tokenizer = LLaMATokenizer.from_pretrained("alpaca_out", cache_dir=CACHE_DIR, local_files_only=True)
def predict(
self,
prompt: str = Input(description=f"Prompt to send to LLaMA."),
n: int = Input(description="Number of output sequences to generate", default=1, ge=1, le=5),
total_tokens: int = Input(
description="Maximum number of tokens for input + generation. A word is generally 2-3 tokens",
ge=1,
default=2000
),
temperature: float = Input(
description="Adjusts randomness of outputs, greater than 1 is random and 0 is deterministic, 0.75 is a good starting value.",
ge=0.01,
le=5,
default=0.75,
),
top_p: float = Input(
description="When decoding text, samples from the top p percentage of most likely tokens; lower to ignore less likely tokens",
ge=0.01,
le=1.0,
default=1.0
),
repetition_penalty: float = Input(
description="Penalty for repeated words in generated text; 1 is no penalty, values greater than 1 discourage repetition, less than 1 encourage it.",
ge=0.01,
le=5,
default=1
)
) -> List[str]:
format_prompt = PROMPT.format_map({'instruction': prompt})
input = self.tokenizer(format_prompt, return_tensors="pt").input_ids.to(self.device)
outputs = self.model.generate(
input,
num_return_sequences=n,
max_length=total_tokens,
do_sample=True,
temperature=temperature,
top_p=top_p,
repetition_penalty=repetition_penalty
)
out = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
# removing prompt b/c it's returned with every input
out = [val.split('Response:')[1] for val in out]
return out
|