Views
No views yet
1Base_prompt = """You are tasked with correcting spelling mistakes in the queries that users submitted to a Persian marketplace.
2
3Output the corrected query in the following JSON format:
4- If the input requires correction, use:
5 {"correction": "<corrected version of the query>"}
6- If the input is correct, use:
7 {"correction": ""}
8Here are some examples:
9"query": "ندل چسبی زنانه" Your answer: {"correction": "صندل چسبی زنانه"}
10"query": "بادکنک جشن تواد" Your answer: {"correction": "بادکنک جشن تولد"}
11"query": "صندلی بادی" Your answer: {"correction": ""}\n"""1//output structring
2def extract_json(text):
3 try:
4 correction = None
5 pos = 0
6 decoder = json.JSONDecoder()
7 while pos < len(text):
8 match = text.find('{"correction":', pos)
9 if match == -1:
10 break
11 try:
12 result, index = decoder.raw_decode(text[match:])
13 correction = result.get('correction')
14 if correction:
15 return correction
16 pos = match + index
17 except json.JSONDecodeError:
18 pos = match + 1
19 return correction
20 except Exception as e:
21 return text
22
23
24//Load Model
25BASE_MODEL_PATH = "meta-llama/Meta-Llama-3.1-8B-Instruct"
26model_name_or_path = "mfsadi/Llama-3.1-8B-spelling-fa"
27base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL_PATH, return_dict=True)
28spelling_model = PeftModel.from_pretrained(base_model, model_name_or_path)
29tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_PATH)
30
31
32//Inference. You need to pass "query".
33prompt = f"""### Human: {spell_checking_prompt} query: {query}\n ### Assistant:"""
34batch = tokenizer(str([prompt]), return_tensors='pt')
35prompt_length = len(batch.get('input_ids')[0])
36max_new_tokens = 50
37with torch.no_grad():
38 output_tokens = spelling_model.generate(**batch.to(device), max_new_tokens=max_new_tokens,
39 repetition_penalty=1.1,
40 do_sample=True,
41 num_beams=2,
42 temperature=0.1,
43 top_k=10,
44 top_p=.5,
45 length_penalty=-1
46 )
47 output = tokenizer.decode(output_tokens[0][prompt_length:], skip_special_tokens=True)
48 return extract_json(output)