outlines

1
from enum import Enum
2

3
from llama_cpp import Llama, LogitsProcessorList
4
from pydantic import BaseModel, constr
5

6
from outlines.generate.processors import JSONLogitsProcessor
7
from outlines.models.llamacpp import LlamaCppTokenizer
8

9

10
class Weapon(str, Enum):
11
    sword = "sword"
12
    axe = "axe"
13
    mace = "mace"
14
    spear = "spear"
15
    bow = "bow"
16
    crossbow = "crossbow"
17

18

19
class Armor(str, Enum):
20
    leather = "leather"
21
    chainmail = "chainmail"
22
    plate = "plate"
23

24

25
class Character(BaseModel):
26
    name: constr(max_length=10)
27
    age: int
28
    armor: Armor
29
    weapon: Weapon
30
    strength: int
31

32

33
if __name__ == "__main__":
34
    llama = Llama("./phi-2.Q4_K_M.gguf")
35
    tokenizer = LlamaCppTokenizer(llama)
36

37
    prompt = "Instruct: You are a leading role play gamer. You have seen thousands of different characters and their attributes.\nPlease return a JSON object with common attributes of an RPG character. Give me a character description\nOutput:"
38

39
    logits_processor = JSONLogitsProcessor(Character, tokenizer)
40

41
    json_str = llama.create_completion(
42
        prompt,
43
        top_k=40,
44
        top_p=0.95,
45
        temperature=0.7,
46
        max_tokens=100,
47
        logits_processor=LogitsProcessorList([logits_processor]),
48
    )["choices"][0]["text"]
49

50
    print(json_str)
51