File size: 2,991 Bytes
7bf4b88 |
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 |
import sys
import os
import torch.nn as nn
import ast
from unsloth import FastLanguageModel
from transformers import TextStreamer
import contractions
import re
from Planning.utils import remove_inner_single_quotes
class Planner(nn.Module):
def __init__(self, dataset_name):
super(Planner, self).__init__()
self.dataset_name = dataset_name
self.checkpoint_path = f"Planning/checkpoints/{dataset_name}/lora_model/"
self.max_seq_length = 2048
self.dtype = None
self.load_in_4bit = True
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = self.checkpoint_path,
max_seq_length = self.max_seq_length,
dtype = self.dtype,
load_in_4bit = self.load_in_4bit
)
FastLanguageModel.for_inference(model)
self.model = model
self.tokenizer = tokenizer
def forward(self, query):
message = {'content': query, 'role': 'user'}
inputs = self.tokenizer.apply_chat_template(
[message],
tokenize = True,
add_generation_prompt = True, # Must add for generation
return_tensors = "pt",
).to("cuda")
text_streamer = TextStreamer(self.tokenizer, skip_prompt = True)
outputs = self.model.generate(input_ids = inputs, streamer = text_streamer, max_new_tokens = 128, # max_new_tokens is the maximum number of new tokens generated beyond the input
use_cache = True, temperature = 1.5, min_p = 0.1) # min_p is a cumulative probability, which makes the generation more diverse
outputs = self.tokenizer.batch_decode(outputs)
parts = outputs[0].split("<|start_header_id|>assistant<|end_header_id|>\n\n")
if len(parts) > 1:
results = parts[1].replace("<|eot_id|>", "")
else:
raise ValueError
# ******* special processing for prime dataset
if self.dataset_name == 'prime':
try:
# Parse the string using ast.literal_eval
parsed_dict = ast.literal_eval(results)
return parsed_dict
except (SyntaxError, ValueError) as e:
print(f"Error parsing the string: {e}")
return {
"Metapath": "",
"Restriction": {}
}
results = contractions.fix(results)
try:
results = ast.literal_eval(results)
except:
print(f"Fail")
try:
results = re.sub(r"\['(.*?)'", remove_inner_single_quotes, results) # TODO: need optimize
results = ast.literal_eval(results)
except:
results = {
"Metapath": "",
"Restriction": {},
}
rg = results
return rg |