Spaces:
Sleeping
Sleeping
File size: 4,859 Bytes
1a212f3 eb93814 1a212f3 eb93814 1a212f3 eb93814 1a212f3 eb93814 1a212f3 eb93814 1a212f3 eb93814 1a212f3 | 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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
import diskcache
import openai
from dotenv import load_dotenv
from inference_utils import sanitize_path_component
load_dotenv()
@dataclass
class LLMResponse:
model: str
config: str
text_a: str
text_b: str
edited_text_a: Optional[str]
response_obj: dict
class AzureClient:
API_KEY_NAME = "AZURE_API_KEY"
def __init__(
self,
model_name: str = "gpt-5.6-terra",
config_name: str = "config.v0.6.json",
cache_directory: Optional[Path] = None,
use_cache: bool = True,
):
self.config_name = config_name
with open(Path(__file__).parent / "configs" / config_name) as f:
self.config = json.load(f)
assert "{text_a}" in self.config["prompt_template"]
assert "{text_b}" in self.config["prompt_template"]
for example in self.config.get("examples", []):
assert "text_a" in example
assert "text_b" in example
assert "edited_text_a" in example
self.model_name = model_name
self.BASE_URL = os.environ.get("AZURE_BASE_URL")
self.cache_directory = cache_directory or Path(__file__).parent / ".llm_cache"
self.client = openai.Client(
api_key=os.environ.get(self.API_KEY_NAME),
base_url=self.BASE_URL,
)
self.cache = None
if use_cache:
self.cache = diskcache.Cache(self.model_cache_dir)
@property
def model_cache_dir(self) -> Path:
sanitized_model_name = sanitize_path_component(self.model_name)
config_stem = sanitize_path_component(self.config_name.removesuffix(".json"))
return self.cache_directory / f"{sanitized_model_name}_{config_stem}"
def _get_cache_key(self, text_a: str, text_b: str) -> tuple[str, str, str, str]:
return (self.model_name, self.config_name, text_a, text_b)
def _completion_extra_kwargs(self) -> dict:
kwargs = {"reasoning_effort": "none"}
for key in ("seed", "temperature"):
if key in self.config:
kwargs[key] = self.config[key]
return kwargs
def _format_prompt(self, text_a: str, text_b: str) -> str:
return self.config["prompt_template"].format(
text_a=text_a,
text_b=text_b,
)
def _build_messages(self, text_a: str, text_b: str) -> list[dict]:
messages = []
if self.config.get("use_examples", False):
for example in self.config.get("examples", []):
messages.append({
"role": "user",
"content": self._format_prompt(example["text_a"], example["text_b"]),
})
messages.append({
"role": "assistant",
"content": example["edited_text_a"],
})
messages.append({
"role": "user",
"content": self._format_prompt(text_a, text_b),
})
return messages
def query(self, text_a: str, text_b: str) -> LLMResponse:
cache_key = self._get_cache_key(text_a, text_b)
if self.cache is not None and cache_key in self.cache:
return self._build_response(text_a, text_b, self.cache[cache_key])
completion = self.client.chat.completions.create(
model=self.model_name,
messages=self._build_messages(text_a, text_b),
**self._completion_extra_kwargs(),
)
response_obj = completion.to_dict()
if self.cache is not None:
self.cache[cache_key] = response_obj
return self._build_response(text_a, text_b, response_obj)
@staticmethod
def _normalize_edited_text(content: str, text_a: str) -> str:
stripped = content.strip()
candidate = stripped
if len(candidate) >= 2 and candidate[0] == candidate[-1] and candidate[0] in "\"'":
candidate = candidate[1:-1].strip()
if candidate.casefold() == "pass":
return text_a
return stripped
def _build_response(self, text_a: str, text_b: str, response_obj: dict) -> LLMResponse:
llm_response = LLMResponse(
model=self.model_name,
config=self.config_name,
text_a=text_a,
text_b=text_b,
edited_text_a=None,
response_obj=response_obj,
)
try:
content = response_obj["choices"][0]["message"]["content"]
except (KeyError, IndexError, TypeError):
return llm_response
if content is None:
return llm_response
llm_response.edited_text_a = self._normalize_edited_text(content, text_a)
return llm_response
|