add support for running title generator only for part of a file

This commit is contained in:
Ashik K 2023-04-06 20:56:18 +02:00
parent 343a7b5fe7
commit c714573520
1 changed files with 15 additions and 7 deletions

View File

@ -3,23 +3,31 @@ import csv
import openai import openai
# check if the user supplied the correct number of command line arguments # check if the user supplied the correct number of command line arguments
if len(sys.argv) != 4: if len(sys.argv) != 6:
print("Usage: python titlegen.py inputfilename.csv outputfilename.csv API_key") print("Usage: python titlegen.py inputfilename.csv startlineno endlineno outputfilename.csv API_key")
sys.exit(1) sys.exit(1)
# set the api key from the command line argument # set the api key from the command line argument
openai.api_key = sys.argv[3] openai.api_key = sys.argv[5]
startlineno = int(sys.argv[2])
endlineno = int(sys.argv[3])
lineno = 0
with open(sys.argv[1], "r") as file: with open(sys.argv[1], "r") as file:
outfile = open(sys.argv[2], "w") outfile = open(sys.argv[4], "w")
writer = csv.writer(outfile) writer = csv.writer(outfile)
reader = csv.reader(file) reader = csv.reader(file)
for row in reader: for row in reader:
print(row[4]) lineno += 1
if (lineno < startlineno or lineno > endlineno):
continue
# summarize the text in the 5th column using OpenAi's GPT-3 # summarize the text in the 5th column using OpenAi's GPT-3
# create a variable called prompt and set it as the concatenatenation of the string "Summarize this # create a variable called prompt and set it as the concatenatenation of the string "Summarize this
# in one sentence:" and the text in the 5th column # in one sentence:" and the text in the 5th column
prompt = "Summarize the following description of a museum item in one English phrase of strictly less than 5 words:" + row[4] desc = row[4]
# extract the first 12 words of the description
trimmed_desc = " ".join(desc.split()[:12])
prompt = "Summarize the following description of a museum item in one English phrase of strictly less than 4 words: " + trimmed_desc
completion = openai.ChatCompletion.create( # 1. Change the function Completion to ChatCompletion completion = openai.ChatCompletion.create(
model = 'gpt-3.5-turbo', model = 'gpt-3.5-turbo',
messages = [ messages = [
{'role': 'user', 'content': prompt} {'role': 'user', 'content': prompt}