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