File size: 2,085 Bytes
c95100d 7902a6f c95100d 7902a6f c95100d 7902a6f c95100d 05d1e4e c95100d |
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 |
import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
peft_model_id = f"telmo000/bloom-positive-reframing"
config = PeftConfig.from_pretrained(peft_model_id)
model = AutoModelForCausalLM.from_pretrained(
config.base_model_name_or_path,
return_dict=True,
load_in_8bit=True,
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
# Load the Lora model
model = PeftModel.from_pretrained(model, peft_model_id)
def make_inference(original_text):
str_strategy_prompt = f"### Negative sentence:\n{original_text}\n\n### Reframing strategy:\n"
batch_1 = tokenizer(
str_strategy_prompt,
return_tensors="pt",
)
with torch.cuda.amp.autocast():
output_tokens_1 = model.generate(**batch_1, max_new_tokens=50)
output_1 = tokenizer.decode(output_tokens_1[0], skip_special_tokens=True)
reframing_strategy = output_1[len(str_strategy_prompt):].partition('\n')[0]
str_reframing_prompt = f"### Negative sentence:\n{original_text}\n\n### Reframing strategy:\n{reframing_strategy}\n\n### Reframing sentence:\n"
batch_2 = tokenizer(
str_reframing_prompt,
return_tensors="pt",
)
with torch.cuda.amp.autocast():
output_tokens_2 = model.generate(**batch_2, max_new_tokens=100)
output_2 = tokenizer.decode(output_tokens_2[0], skip_special_tokens=True)
reframing_sentence = output_2[len(str_reframing_prompt):]
return reframing_sentence
if __name__ == "__main__":
# make a gradio interface
import gradio as gr
gr.Interface(
make_inference,
[
gr.inputs.Textbox(lines=3, label="Original Text"),
],
gr.outputs.Textbox(label="Reframed Text"),
title="Bloom positive reframing",
description="Bloom positive reframing is a BLOOM-base generative model adjusted to the sentiment transfer task, where the objective is to reverse the sentiment polarity of a text without contradicting the original meaning. ",
).launch() |