Spaces:
Running
Running
File size: 4,086 Bytes
eb60b6b | 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 | from __future__ import annotations
import time
from dataclasses import dataclass
from typing import Any
from openai import AzureOpenAI, OpenAI
from longmemeval.config import Settings
from longmemeval.prompts import build_answer_messages
from longmemeval.sota import majority_vote
@dataclass
class AnswerResult:
answer: str
model: str
latency_ms: float
prompt_tokens: int | None
completion_tokens: int | None
total_tokens: int | None
sc_samples_used: int = 1
def _make_openai_client(settings: Settings) -> OpenAI | AzureOpenAI:
settings.require_llm()
if settings.answer_azure_endpoint:
return AzureOpenAI(
api_key=settings.openai_api_key,
azure_endpoint=settings.answer_azure_endpoint,
api_version=settings.answer_azure_api_version,
)
kwargs: dict[str, Any] = {"api_key": settings.openai_api_key}
if settings.llm_base_url:
kwargs["base_url"] = settings.llm_base_url
return OpenAI(**kwargs)
def _one_completion(
client: OpenAI | AzureOpenAI,
*,
model_name: str,
messages: list[dict[str, str]],
temperature: float,
) -> tuple[str, Any]:
response = client.chat.completions.create(
model=model_name,
messages=messages,
temperature=temperature,
)
content = (response.choices[0].message.content or "").strip()
return content, response
def answer_question(
settings: Settings,
question: str,
context: dict[str, Any],
*,
question_date: str | None = None,
model: str | None = None,
sc_samples: int | None = None,
sc_temperature: float | None = None,
) -> AnswerResult:
client = _make_openai_client(settings)
model_name = model or settings.answer_model
messages = build_answer_messages(
question, context, question_date=question_date
)
n = max(1, int(sc_samples if sc_samples is not None else settings.sc_samples))
temperature = (
float(sc_temperature)
if sc_temperature is not None
else (settings.sc_temperature if n > 1 else 0.0)
)
started = time.perf_counter()
prompt_tokens = 0
completion_tokens = 0
total_tokens = 0
samples: list[str] = []
last_response: Any = None
for _ in range(n):
content, response = _one_completion(
client,
model_name=model_name,
messages=messages,
temperature=0.0 if n == 1 else temperature,
)
last_response = response
if response.usage:
prompt_tokens += getattr(response.usage, "prompt_tokens", 0) or 0
completion_tokens += getattr(response.usage, "completion_tokens", 0) or 0
total_tokens += getattr(response.usage, "total_tokens", 0) or 0
if content:
samples.append(content)
content = majority_vote(samples) if samples else ""
if not content:
retry_messages = messages + [
{
"role": "user",
"content": (
"Your previous reply was empty. Answer the question now using the "
"retrieved context. Prefer a short factual list or phrase."
),
}
]
content, response = _one_completion(
client,
model_name=model_name,
messages=retry_messages,
temperature=0,
)
last_response = response
if response.usage:
prompt_tokens += getattr(response.usage, "prompt_tokens", 0) or 0
completion_tokens += getattr(response.usage, "completion_tokens", 0) or 0
total_tokens += getattr(response.usage, "total_tokens", 0) or 0
latency_ms = (time.perf_counter() - started) * 1000
return AnswerResult(
answer=content,
model=getattr(last_response, "model", None) or model_name,
latency_ms=latency_ms,
prompt_tokens=prompt_tokens or None,
completion_tokens=completion_tokens or None,
total_tokens=total_tokens or None,
sc_samples_used=n,
)
|