Unverified Commit 185ad31f authored by Eric's avatar Eric Committed by GitHub
Browse files

[Bugfix] use diskcache in outlines _get_guide #5436 (#6203)

parent 543aa485
......@@ -21,6 +21,7 @@ from functools import lru_cache
from typing import Callable, DefaultDict, Dict, List, Union
import torch
from outlines.caching import cache
from outlines.fsm.guide import CFGGuide, Generate, Guide, RegexGuide, Write
from outlines.fsm.json_schema import build_regex_from_schema
from pydantic import BaseModel
......@@ -67,7 +68,7 @@ class BaseLogitsProcessor:
class RegexLogitsProcessor(BaseLogitsProcessor):
@classmethod
@lru_cache(maxsize=32)
@cache()
def _get_guide(cls, regex_string: str,
tokenizer: PreTrainedTokenizerBase) -> Guide:
tokenizer = _adapt_tokenizer(tokenizer)
......@@ -126,7 +127,7 @@ class JSONLogitsProcessor(RegexLogitsProcessor):
class CFGLogitsProcessor(BaseLogitsProcessor):
@classmethod
@lru_cache(maxsize=32)
@cache()
def _get_guide(cls, cfg: str, tokenizer: PreTrainedTokenizerBase) -> Guide:
tokenizer = _adapt_tokenizer(tokenizer)
return CFGGuide(cfg, tokenizer)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment