File size: 1,705 Bytes
27757e7 d93389c 27757e7 |
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 |
import json
import random
random.seed(999)
import torch
from transformers import Qwen2ForSequenceClassification, AutoTokenizer
import gradio as gr
from datetime import datetime
torch.set_grad_enabled(False)
model = Qwen2ForSequenceClassification.from_pretrained("Thouph/prompt2tag-qwen2-0.5b-v0.1", num_labels = 9940, device_map="cpu")
model.eval()
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B")
with open("tags_9940.json", "r") as file:
allowed_tags = json.load(file)
allowed_tags = sorted(allowed_tags)
def create_tags(prompt, threshold):
inputs = tokenizer(
prompt,
padding="do_not_pad",
max_length=512,
truncation=True,
return_tensors="pt",
)
output = model(**inputs).logits
output = torch.nn.functional.sigmoid(output)
indices = torch.where(output > threshold)
values = output[indices]
indices = indices[1]
values = values.squeeze()
temp = []
tag_score = dict()
for i in range(indices.size(0)):
temp.append([allowed_tags[indices[i]], values[i].item()])
tag_score[allowed_tags[indices[i]]] = values[i].item()
temp = [t[0] for t in temp]
text_no_impl = " ".join(temp)
current_datetime = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
print(f"{current_datetime}: finished.")
return text_no_impl, tag_score
demo = gr.Interface(
create_tags,
inputs=[
gr.TextArea(label="Prompt",),
gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.40, label="Threshold")
],
outputs=[
gr.Textbox(label="Tag String"),
gr.Label(label="Tag Predictions", num_top_classes=200),
],
allow_flagging="never",
)
demo.launch()
|