import time import numpy as np import torch from transformers import AutoTokenizer, AutoModelForCausalLM, AutoProcessor from transformers import Gemma3ForConditionalGeneration import langdetect from langdetect import detect from langdetect import DetectorFactory from pathlib import Path import datetime from core.prompts import prompts DetectorFactory.seed = 0 available_models = {"LLAMA1B": "meta-llama/Llama-3.2-1B-Instruct", "GEMMA2B": "google/gemma-2-2b-it", "SALAMANDRA2B": "BSC-LT/salamandra-2b-instruct", "LLAMA3B": "meta-llama/Llama-3.2-3B-Instruct", "GEMMA4B": "google/gemma-3-4b-it", "OLMO7B": "allenai/OLMo-2-1124-7B-Instruct", "SALAMANDRA7B": "BSC-LT/salamandra-7b-instruct", "LLAMA8B": "meta-llama/Llama-3.1-8B-Instruct", "GEMMA9B": "google/gemma-2-9b-it", "GEMMA12B": "google/gemma-3-12b-it", "GEMMA27B": "google/gemma-3-27b-it"} # Inspired with https://huggingface.co/google/gemma-3-12b-it # More info in https://huggingface.co/docs/transformers/main/en/model_doc/gemma3 class Correctifier(object): def __init__(self, selected_model, access_token, device, lang="en"): self.lang = lang pretrained_model = available_models[selected_model] self.pretrained_model = pretrained_model self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model, token=access_token) self.processor = AutoProcessor.from_pretrained(pretrained_model, token=access_token) # If CUDA device is chosen and more than one available, use all device_map = "auto" if device.type == 'cuda' and torch.cuda.is_available() and torch.cuda.device_count() > 1 else device print("Using {} device".format(device_map), device.type, torch.cuda.is_available(), torch.cuda.device_count()) if pretrained_model.startswith('google/gemma-3'): self.model = Gemma3ForConditionalGeneration.from_pretrained(pretrained_model, token=access_token, torch_dtype=torch.bfloat16, device_map=device_map).eval() elif pretrained_model.startswith('google/gemma-2'): self.model = Gemma3ForConditionalGeneration.from_pretrained(pretrained_model, token=access_token, torch_dtype=torch.float16, device_map=device_map).eval() else: self.model = AutoModelForCausalLM.from_pretrained(pretrained_model, torch_dtype=torch.bfloat16, token=access_token, device_map=device_map).eval() self.model.generation_config.pad_token_id = self.tokenizer.eos_token_id self.device = device print("Model loaded.") if torch.cuda.is_available(): for i in range(torch.cuda.device_count()): print("Device " + str(i) + ' : ' + torch.cuda.get_device_properties(i).name) print("Total memory: " + str(torch.cuda.get_device_properties(i).total_memory)) print("Reserved: " + str(torch.cuda.memory_reserved(i))) print("Allocated: " + str(torch.cuda.memory_allocated(i))) print("Model on GPU: " + str(round(100 * np.sum([param.is_cuda for param in self.model.parameters()]) / len( list(self.model.parameters())))) + "%") def correct(self, sentence, force_prompt=None, force_raw_output=False): # print("Simplifying:\n" + sentence) new_tokens = 1000 # Prompt based on https://aclanthology.org/2023.emnlp-main.821/ if force_prompt is not None: prompt = force_prompt else: prompt = prompts[self.lang] prompt += "\nINPUT:\n" + sentence # + "\nOUTPUT:\n" messages = [ # {"role": "system","content": [{"type": "text", "text": "You are a helpful assistant."}] }, { "role": "user", "content": [{"type": "text", "text": prompt}] } ] if self.pretrained_model.startswith("allenai/OLMo"): inputs = self.tokenizer([prompt], return_tensors='pt', return_token_type_ids=False) elif self.pretrained_model.startswith("BSC-LT/salamandra"): prompt_w_templ = [{"role": "user", "content": prompt}] messages = self.tokenizer.apply_chat_template( prompt_w_templ, tokenize=False, add_generation_prompt=True, ) inputs = self.tokenizer.encode(messages, add_special_tokens=False, return_tensors="pt") else: inputs = self.processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt") if self.pretrained_model.startswith("google/gemma-3"): inputs = inputs.to(self.device, dtype=torch.bfloat16) elif self.pretrained_model.startswith("google/gemma-2"): inputs = inputs.to(self.device, dtype=torch.float16) else: inputs = inputs.to(self.device) if self.pretrained_model.startswith("BSC-LT/salamandra"): input_len = 0 print(inputs.shape) input_len = inputs.shape[-1] else: input_len = inputs["input_ids"].shape[-1] start = time.time() if self.pretrained_model.startswith("BSC-LT/salamandra"): with torch.inference_mode(): # generation = self.model.generate(input_ids=inputs.to(self.device), max_new_tokens=200) generation = self.model.generate(input_ids=inputs, max_new_tokens=200, do_sample=False, top_p=None, temperature=None) generation = generation[0][input_len:] else: with torch.inference_mode(): # Using greedy decoding, see others in https://huggingface.co/docs/transformers/en/generation_strategies generation = self.model.generate(**inputs, max_new_tokens=new_tokens, do_sample=False, top_p=None, temperature=None) generation = generation[0][input_len:] end = time.time() evaluate_time = end - start print("Response time: " + str(evaluate_time)) output_text = self.processor.decode(generation, skip_special_tokens=True) # print(output_text) if force_raw_output: return output_text else: return self.parse_responses(sentence, output_text) def parse_responses(self, sentence, output_text): response = output_text # [output_text.find('OUTPUT:\n'):] # for end_marker in ["", "<|eot_id|>", "```<|end_of_text|>", "", "<|endoftext|>"]: # if response.endswith(end_marker): # response = response[0:(-(len(end_marker)))] responses = response.split("\n") responses = [response for response in responses if not ( response.strip() in ['', 'OUTPUT:', 'INPUT:', sentence] or response.startswith( 'Rephrasing ') or response.startswith('Here are ') or response.startswith('Rewrite '))] responses = [response.lstrip('1234567890-').lstrip('.').lstrip() for response in responses] if (Path(langdetect.__file__).parents[0] / 'profiles' / self.lang).exists(): responses = [response for response in responses if len(response) < 10 or detect(response) == self.lang] if len(responses) > 0: return responses[0] else: return ''