add: output response to csv

qqpann · Jul 22, 2023 · f2fa1b0 · f2fa1b0
1 parent ac73e14
commit f2fa1b0
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 3 deletions.
diff --git a/.gitignore b/.gitignore
@@ -160,4 +160,5 @@ cython_debug/
 #.idea/
 
 models/*
+out/*
 !.gitkeep
diff --git a/out/.gitkeep b/out/.gitkeep
diff --git a/requirements.txt b/requirements.txt
@@ -5,3 +5,4 @@ mypy
 openai
 pylint
 python-dotenv
+pandas
diff --git a/run.py b/run.py
@@ -1,15 +1,31 @@
+import pandas as pd
+from tqdm import tqdm
+
 from model import chat
 
 DEBUG = True
 
 
 def main():
+    df = pd.DataFrame(columns=["question", "ai_answer", "modified_answer"])
     with open("data/aituber_question_dataset/question.txt", "r") as f:
-        for line in f.readlines():
-            print("question >", line)
-            print("ai answer >", chat(line))
+        for line in tqdm(f.readlines()):
+            line = line.strip()
+            question = line
+            ai_answer = chat(line)
+            modified_answer = ""
+            df = df._append(
+                {
+                    "question": question,
+                    "ai_answer": ai_answer,
+                    "modified_answer": modified_answer,
+                },
+                ignore_index=True,
+            )
+
             if DEBUG:
                 break
+    df.to_csv(f"out/question_answer.csv", index=False)
 
 
 if __name__ == "__main__":
-Original file line number
+Diff line change
@@ Expand Up / @@ -160,4 +160,5 @@ cython_debug/ @@
     #.idea/
     models/*
+    out/*
     !.gitkeep