mod test
This commit is contained in:
parent
6816562163
commit
2ea7c63c7f
13
test.py
13
test.py
|
|
@ -3,13 +3,20 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
|
|||
model_path = "./trained_model/gpt"
|
||||
model = AutoModelForCausalLM.from_pretrained(model_path)
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
model.config.pad_token_id = tokenizer.eos_token_id
|
||||
|
||||
def generate_response(prompt, max_length=100):
|
||||
inputs = tokenizer(prompt, return_tensors="pt")
|
||||
outputs = model.generate(inputs.input_ids, max_length=max_length, num_return_sequences=1, do_sample=True)
|
||||
inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
|
||||
outputs = model.generate(
|
||||
inputs.input_ids,
|
||||
attention_mask=inputs.attention_mask,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
max_length=100
|
||||
)
|
||||
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
|
||||
return response
|
||||
|
||||
prompt = "Jakie są prawa pracownika zgodnie z Kodeksem pracy?"
|
||||
prompt = "Ile dni urlopu przysługuje pracownikowi z 5 letnim stazem pracy w pełnym wymiarze pracy?"
|
||||
response = generate_response(prompt)
|
||||
print(response)
|
||||
Loading…
Reference in New Issue