Add gpt2 solution
This commit is contained in:
parent
a7ec11ca27
commit
f8892f9209
18
gpt_predict.py
Normal file
18
gpt_predict.py
Normal file
@ -0,0 +1,18 @@
|
|||||||
|
from transformers import pipeline
|
||||||
|
import lzma
|
||||||
|
|
||||||
|
generator = pipeline("text-generation", model="gpt2")
|
||||||
|
|
||||||
|
with open("test-A/in.tsv", "r") as input_file, open(
|
||||||
|
"test-A/out.tsv", "w"
|
||||||
|
) as output_file:
|
||||||
|
for line in input_file:
|
||||||
|
line = line.rstrip()
|
||||||
|
line = line.replace("\\n", " ")
|
||||||
|
|
||||||
|
prompt = line.split("\t")[6]
|
||||||
|
|
||||||
|
result = generator(prompt, max_new_tokens=1, num_return_sequences=1)[0][
|
||||||
|
"generated_text"
|
||||||
|
]
|
||||||
|
output_file.write(f"{result.split()[-1]}:1\n")
|
7414
test-A/out.tsv
Normal file
7414
test-A/out.tsv
Normal file
File diff suppressed because it is too large
Load Diff
Loading…
Reference in New Issue
Block a user