Spaces:
Runtime error
Runtime error
File size: 3,239 Bytes
d8e827d ac51a1c d8e827d |
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 |
import logging
import re
import string
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
GEMMA_WORD_PATTERNS = [
"(?<=\*)(.*?)(?=\*)",
'(?<=")(.*?)(?=")',
]
def query_hf(
query: str,
model: AutoModelForCausalLM,
tokenizer: AutoTokenizer,
generation_config: dict,
device: str,
) -> str:
"""Queries an LLM model using the Vertex AI API.
Args:
query (str): Query sent to the Vertex API
model (str): Model target by Vertex
generation_config (dict): Configurations used by the model
Returns:
str: Vertex AI text response
"""
generation_config = GenerationConfig(
do_sample=True,
max_new_tokens=generation_config["max_output_tokens"],
top_k=generation_config["top_k"],
top_p=generation_config["top_p"],
temperature=generation_config["temperature"],
)
input_ids = tokenizer(query, return_tensors="pt").to(device)
outputs = model.generate(**input_ids, generation_config=generation_config)
outputs = tokenizer.decode(outputs[0], skip_special_tokens=True)
outputs = outputs.replace(query, "")
return outputs
def query_word(
category: str,
model: AutoModelForCausalLM,
tokenizer: AutoTokenizer,
generation_config: dict,
device: str,
) -> str:
"""Queries a word to be used for the hangman game.
Args:
category (str): Category used as source sample a word
model (str): Model target by Vertex
generation_config (dict): Configurations used by the model
Returns:
str: Queried word
"""
logger.info(f"Quering word for category: '{category}'...")
query = f"Name a single existing {category}."
matched_word = ""
while not matched_word:
word = query_hf(query, model, tokenizer, generation_config, device)
logger.info(f"Evaluating result: '{word}'...")
# Extract word of interest from Gemma's output
for pattern in GEMMA_WORD_PATTERNS:
matched_words = re.findall(rf"{pattern}", word)
matched_words = [x for x in matched_words if x != ""]
if matched_words:
matched_word = matched_words[-1]
matched_word = matched_word.translate(str.maketrans("", "", string.punctuation))
matched_word = matched_word.lower()
logger.info("Word queried successful")
return matched_word
def query_hint(
word: str,
model: AutoModelForCausalLM,
tokenizer: AutoTokenizer,
generation_config: dict,
device: str,
) -> str:
"""Queries a hint for the hangman game.
Args:
word (str): Word used as source to create the hint
model (str): Model target by Vertex
generation_config (dict): Configurations used by the model
Returns:
str: Queried hint
"""
logger.info(f"Quering hint for word: '{word}'...")
query = f"Describe the word '{word}' without mentioning it."
hint = query_hf(query, model, tokenizer, generation_config, device)
hint = re.sub(re.escape(word), "***", hint, flags=re.IGNORECASE)
logger.info("Hint queried successful")
return hint
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__file__)
|