Views
No views yet
transformers library and use it to generate feedback for Korean STT Data.1
2import torch
3from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
4
5# Load the pre-trained model and tokenizer for STT correction
6model = AutoModelForCausalLM.from_pretrained("stt-error-correction-model")
7tokenizer = AutoTokenizer.from_pretrained("stt-error-correction-tokenizer")
8
9# Create a text generation pipeline using the fine-tuned model and tokenizer
10pipe_finetuned = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512)
11
12# Define a placeholder for the STT error text
13stt_error_text = "안뇽하새요"
14
15# Construct a list of messages for STT correction
16messages = [
17 {
18 "role": "user",
19 "content": (
20 f"STT 오류가 포함된 텍스트를 올바르게 수정해주세요.\n"
21 f"STT 오류 텍스트: {stt_error_text}\n"
22 )
23 }
24]
25
26# Prepare the input prompt using the tokenizer
27prompt = pipe_finetuned.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
28
29# Generate correction by passing the formatted prompt to the pipeline
30outputs = pipe_finetuned(
31 prompt,
32 do_sample=True, # Enable sampling to generate diverse outputs
33 temperature=0.2, # Control randomness in text generation (lower value makes the output more focused)
34 top_k=50, # Limit the sampling pool to the top 50 tokens
35 top_p=0.95, # Use nucleus sampling to focus on the top 95% of probability mass
36 add_special_tokens=True # Include special tokens as per the model's requirements
37)
38
39# Print the generated correction
40print(outputs[0]["generated_text"][len(prompt):])
41
42
43
44
45
46