19 lines
510 B
Python
19 lines
510 B
Python
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")
|