Add generators

Author: JKusio

.gitignore

@@ -4,4 +4,5 @@ env
 output
 __pycache__
 .vscode
-openai_batches
+openai_batches
+models

src/common/names.py

@@ -12,6 +12,11 @@
 
 QA_MODEL_NAMES = ["radlab/polish-qa-v2", "timpal0l/mdeberta-v3-base-squad2"]
 
+INST_MODEL_PATHS = [
+    "../../models/Bielik-7B-Instruct-v0.1-q4",
+    "../../models/Mistral-7B-Instruct-v0.2-q4",
+]
+
 OPENAI_EMBEDDING_MODEL_NAMES = ["text-embedding-3-large"]
 
 QUERY_PREFIX_MAP = {

src/generators/generator.py

@@ -0,0 +1,8 @@
+from abc import ABC, abstractmethod
+from common.passage import Passage
+
+
+class Generator(ABC):
+    @abstractmethod
+    def generate_answer(self, query: str, passages: list[Passage]) -> str:
+        pass

src/generators/instruction_generator.py

@@ -0,0 +1,37 @@
+from common.passage import Passage
+from generators.generator import Generator
+from mlx_lm import load, generate
+
+
+class InstructionGenerator(Generator):
+    def __init__(self, model_name: str):
+        self.model_name = model_name
+
+        model, tokenizer = load(model_name)
+
+        self.model = model
+        self.tokenizer = tokenizer
+
+    def generate_answer(self, query: str, passages: list[Passage]) -> str:
+        context = " ".join([passage.context for passage in passages]).replace("\n", " ")
+
+        prompt = f"""
+        [INST]
+        Wygeneruj krótką odpowiedź na pytanie wyłącznie na podstawie poniższego kontekstu:
+        {context}
+
+        Pytanie: {query}
+        [/INST]
+        """
+
+        response = generate(self.model, self.tokenizer, prompt=prompt, max_tokens=200)
+
+        stripped_response = (
+            response.replace("<s>", "")
+            .replace("</s>", "")
+            .replace("[INST]", "")
+            .replace("[/INST]", "")
+            .strip()
+        )
+
+        return stripped_response

src/generators/question_answering_generator.py

@@ -0,0 +1,15 @@
+from common.passage import Passage
+from transformers import pipeline
+from generators.generator import Generator
+
+
+class QuestionAnsweringGenerator(Generator):
+    def __init__(self, model_name: str):
+        self.model_name = model_name
+
+        self.pipeline = pipeline("question-answering", model=model_name)
+
+    def generate_answer(self, query: str, passages: list[Passage]) -> str:
+        context = " ".join([passage.context for passage in passages]).replace("\n", " ")
+
+        return self.pipeline(question=query, context=context)["answer"]

src/notebooks/00_test.ipynb

@@ -11,6 +11,46 @@
     "sys.path.append(\"../\")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from generators.instruction_generator import InstructionGenerator\n",
+    "\n",
+    "generator = InstructionGenerator(\"../../models/Mistral-7B-Instruct-v0.2-q4\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Model Hunter T.7A był wykorzystywany do szkolenia załóg samolotów w Royal Air Force po przebudowie z dodaniem systemu IFIS.\n"
+     ]
+    }
+   ],
+   "source": [
+    "from common.passage import Passage\n",
+    "\n",
+    "\n",
+    "query = \"W jakim sposób eksploatowano model T.7A?\"\n",
+    "\n",
+    "context = \"Z 55 zamówionych Hunterów T.7 45 trafiło do RAF-u, a pozostałych 10 do Royal Navy. Ponadto RAF dostał sześć maszyn przebudowanych z myśliwskich Hunterów F.4. Pięć maszyn zmodyfikowano przez dodanie systemu IFIS (OR946 Integrated Flight Instrumentation System) stosowanego na pokładach samolotów English Electric Lightning i Blackburn Buccaneer. Samoloty te zyskały nowe oznaczenie Hunter T.7A i były wykorzystywane do szkolenia załóg samolotów. Kolejna wersja Huntera oznaczona jako Hunter T.8 została pozbawiona części awioniki, ale wyposażona w hak do lądowania na lotniskowcach. Maszyny te nie były używane na morzu, a jedynie do ćwiczeń lądowania na lądzie. W 1957 opracowano kolejną wersję samolotu napędzaną nowszym i mocniejszym silnikiem Avon serii 200. Większa moc jednostki napędowej pozwoliła na zainstalowanie dwóch działek ADEN, ale tak wyposażona maszyna oznaczona jako Hunter T.66 nigdy nie weszła do służby w brytyjskich siłach zbrojnych. Jedna z maszyn tego typu oznaczona jako Hunter T.12 była testowana z awioniką ówcześnie projektowanego samolotu BAC TSR.2. Po zamknięciu projektu TSR.2 służyła ona w Royal Aircraft Establishment w Farnborough jako samolot testowy, między innymi do prób systemu fly by wire.\t\"\n",
+    "\n",
+    "passage = Passage('', '', context, 0, '', '')\n",
+    "\n",
+    "\n",
+    "response = generator.generate_answer(query, [passage])\n",
+    "\n",
+    "print(response)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 3,