Spaces:
Running
Running
File size: 3,336 Bytes
23df63d |
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 |
"""Utilities for the Fujitsu-LLM-KG-8x7B models.
"""
from typing import Literal, Sequence, Tuple
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
###############################################################################
# Generation
###############################################################################
class Fujitsu_LLM_KG:
"""The Fujitsu-LLM-KG-8x7B model.
"""
def __init__(self, model_id: str, *, device_map: str = "auto") -> None:
"""Initializes the model and tokenizer.
"""
self.model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map=device_map,
torch_dtype=torch.bfloat16,
low_cpu_mem_usage=True,
)
self.tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left")
self.tokenizer.pad_token = self.tokenizer.eos_token
def generate(self, prompt:str,
*,
max_new_tokens: int = 2048,
num_beams: int = 1,
) -> str:
"""Generate an answer.
"""
tokenized = self.tokenizer(prompt, return_tensors="pt", padding=True)
with torch.no_grad():
outputs = self.model.generate(
tokenized["input_ids"].to("cuda"),
attention_mask=tokenized["attention_mask"].to("cuda"),
pad_token_id=self.tokenizer.eos_token_id,
max_new_tokens=max_new_tokens,
do_sample=False,
num_beams=num_beams,
)
answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)[len(prompt):]
return answer
###############################################################################
# Extraction
###############################################################################
def extract_turtle(text: str, *, with_rationale = False) -> str:
"""Extracts the RDF Turtle part from the output text of Fujitsu-LLM-KG-8x7B_inst-infer model.
"""
TOKENS = ["<", "rel:", "rdf:", "]"]
if with_rationale:
TOKENS.append("#@")
turtle = ""
for line in text.splitlines():
line_ = line.strip()
if line == "" or any(line_.startswith(c) for c in TOKENS):
if turtle:
turtle += "\n"
turtle += line
return turtle
def extract_answer(text: str) -> Tuple[str, Sequence[str]]:
"""Extracts the final answer part from the output text of Fujitsu-LLM-KG-8x7B_inst-infer model.
"""
path = []
answer = ""
state: Literal["path", "answer"] = "path"
for line in text.splitlines():
if line.strip() and "```" not in line and "## " not in line:
if state == "path":
path.append(line)
elif state == "answer":
if answer:
answer += "\n"
answer += line
if "## Explore Path" in line:
state = "path"
path = []
elif "## Answer" in line:
state = "answer"
answer = ""
elif "```" in line and answer:
break
path = tuple(p.strip() for p in path)
answer = answer.strip()
return answer, path
|