v6

Untitled.ipynb

@@ -0,0 +1,82 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "58e6167e-1e30-4bbf-93a1-af6ab5cb4673",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "\n",
+    "def preprocess_test(test_solutions_path: str, test_tasks_path: str, test_tests_path: str, save_path: str) -> None:\n",
+    "    def read_files(*paths):\n",
+    "        return (pd.read_excel(path) for path in paths)\n",
+    "\n",
+    "    solutions_df, tasks_df, tests_df = read_files(test_solutions_path, test_tasks_path, test_tests_path)\n",
+    "\n",
+    "    test_dataset = pd.merge(pd.merge(solutions_df, tasks_df, left_on='task_id', right_on='id', how='inner'), \n",
+    "                            tests_df, on='task_id', how='inner')\n",
+    "\n",
+    "    test_dataset['input_output'] = test_dataset.apply(\n",
+    "        lambda row: f\"{row['input']}->{row['output']}\" if pd.notna(row['input']) or pd.notna(row['output']) else \"\", \n",
+    "        axis=1\n",
+    "    )\n",
+    "\n",
+    "    test_dataset = test_dataset[['id', 'student_solution', 'description', 'author_solution', 'input_output']]\n",
+    "    test_dataset.to_excel(save_path, index=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "e561a852-007c-4b78-a7e1-697f80272169",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "base = 'data/raw/test/{}'\n",
+    "\n",
+    "preprocess_test(base.format('solutions.xlsx'), base.format('tasks.xlsx'), base.format('tests.xlsx'), './test.xlsx')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "2c15dbb7-1034-41a7-96be-69e2418bf98e",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "d923d6c4-f962-4646-a041-999862bb3950",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "preprocess_test(base.format('solutions.xlsx'), base.format('tasks.xlsx'), base.format('tests.xlsx'), './test2.xlsx')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python [conda env:.conda-poetry]",
+   "language": "python",
+   "name": "conda-env-.conda-poetry-py"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

app/models/qwen.py

@@ -37,7 +37,7 @@ class Qwen(BaseModel):
         self.max_tokens = max_tokens
         self.messages = []
 
-    def ask(self, user_message: str, clear_history: bool = True) -> Optional[str]:
+    def ask(self, user_message: str, clear_history: bool = True, debug: bool = False) -> Optional[str]:
         if clear_history:
             self.messages = []
             if self.system_prompt:
@@ -47,6 +47,9 @@ class Qwen(BaseModel):
 
         prompt_text = self._build_prompt()
 
+        if debug:
+            print(prompt_text)
+
         inputs = self.tokenizer(prompt_text, return_tensors="pt").to(self.device)
 
         with torch.no_grad():

app/utils/preprocess.py

@@ -1,9 +1,27 @@
 import pandas as pd
 
-def preprocess_test(test_solutions_path: str, test_tasks_path: str, save_path: str) -> None:
+def preprocess_test(test_solutions_path: str, test_tasks_path: str, test_tests_path: str, save_path: str) -> None:
     solutions_df = pd.read_excel(test_solutions_path)
     tasks_df = pd.read_excel(test_tasks_path)
+    tests_df = pd.read_excel(test_tests_path)
 
-    preprocessed_df = solutions_df.merge(tasks_df[['id', 'description']], left_on='task_id', right_on='id', how='left')
+    preprocessed_df = solutions_df.merge(tasks_df[['id', 'description', 'author_solution']], 
+                                         left_on='task_id', right_on='id', how='left')
 
-    preprocessed_df[['description', 'student_solution']].to_excel(save_path, index=False)
+    preprocessed_df = preprocessed_df.merge(tests_df[['task_id', 'input', 'output']], 
+                                            left_on='task_id', right_on='task_id', how='left')
+
+    preprocessed_df['input_output'] = preprocessed_df.apply(
+        lambda row: f"{row['input']}-{row['output']}" if pd.notna(row['input']) or pd.notna(row['output']) else "", 
+        axis=1
+    )
+
+    grouped_df = preprocessed_df.groupby('id_x').agg({
+        'student_solution': 'first',
+        'description': 'first',
+        'author_solution': 'first',
+        'input_output': lambda x: '\n'.join(filter(None, x))
+    }).reset_index()
+
+    grouped_df = grouped_df.rename(columns={'id_x': 'id'})
+    grouped_df.to_excel(save_path, index=False)

app/utils/submit.py

@@ -40,9 +40,9 @@ def generate_submit(tests_path: str, predict_func: Callable, save_path: str, use
         idx = tests.index[i]
         solution_row = tests.iloc[i]
 
-        input_text = f"{solution_row['description']}\n\n{solution_row['student_solution']}"
+        input_text = f"Вводные данные:\n{solution_row['description']}\n\nКод студента:\n{solution_row['student_solution']}\n\nКод автора:\n{solution_row['author_solution']}\n\nТестовые условия:\n{solution_row['input_output']}"
         text = predict_func(input_text)
 
         embedding = embedding2string(get_sentence_embedding(text))
-        submit_df.loc[i] = [i, text, embedding]
+        submit_df.loc[i] = [solution_row['id'], text, embedding]
     submit_df.to_csv(save_path, index=False)

main.py

@@ -18,7 +18,7 @@ if __name__ == "__main__":
         system_prompt=system_prompt,
     )
 
-    preprocess_test("data/raw/test/solutions.xlsx", "data/raw/test/tasks.xlsx", "data/processed/test.xlsx")
+    preprocess_test("data/raw/test/solutions.xlsx", "data/raw/test/tasks.xlsx", "data/raw/test/tests.xlsx", "data/processed/test.xlsx")
 
     # Predict, ёмаё)
     def predict(input_text: str) -> str: