Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3model_name = "cnmoro/Qwen2.5-0.5B-Portuguese-Prompt-Decomposer"
4tokenizer = AutoTokenizer.from_pretrained(model_name)
5model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda")
6
7def decompose_prompt(prompt, _type="textual"):
8 termos_tipo = "textuais" if _type == "textual" else "semânticos"
9 messages = [
10 {
11 "role": "system",
12 "content": f"Você decompõe o texto do usuário em termos de busca {termos_tipo}, no formato JSON"
13 },
14 {
15 "role": "user",
16 "content": prompt
17 }
18 ]
19
20 gen_input = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
21 inputs = tokenizer(gen_input, return_tensors="pt", return_token_type_ids=False).to("cuda")
22 outputs = model.generate(**inputs, max_new_tokens=384, temperature=0.3, do_sample=True, repetition_penalty=1.2)
23
24 input_length = inputs.input_ids.shape[1]
25 generated_tokens = outputs[0, (input_length):]
26 return tokenizer.decode(generated_tokens, skip_special_tokens=True)
27
28print(
29 decompose_prompt(
30 "Com a edição da Lei n. 12.618/2012, quais foram os principais impactos sobre os servidores públicos, considerando que restaram consolidadas duas situações jurídicas previdenciárias?",
31 "textual"
32 )
33)
34# ```json
35# {
36# "search_terms": {
37# "textual": [
38# "Lei 12618",
39# "servidor público previdência",
40# "impacto lei 12618"
41# ]
42# }
43# }
44# ```
45
46#########################
47
48print(
49 decompose_prompt(
50 "Com a edição da Lei n. 12.618/2012, quais foram os principais impactos sobre os servidores públicos, considerando que restaram consolidadas duas situações jurídicas previdenciárias?",
51 "semantic"
52 )
53)
54# ```json
55# {
56# "search_terms": {
57# "semantic": [
58# "impacto lei servidores",
59# "reforma previdência servidores"
60# ]
61# }
62# }
63# ```