implementation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tuwaiq Academy, SDAIA, ALLAM-Challenge, 2024\n",
    "\n",
    "***Abdulrahman Alshahrani, Abdullah Alwakeel***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "This notebook is our implementation of our idea in the Allam-Challenge. It is essentially a structured approach to using a Language Model (LLM) to \"force\" the LLM to generate poetry in any of the many poetry types (bohours). Through this notebook, we aim to explore and demonstrate the capabilities of LLMs in generating structured poetic forms, adhering to the rules and styles of various traditional poetry types."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Explination\n",
    "\n",
    "This program is designed to generate Arabic poetry, specifically focusing on the correct structure and rules of Arabic poetry, such as maintaining correct Wazn (rhythm) and Qafiya (rhyme). It utilizes machine learning models, validation mechanisms, and feedback loops to ensure the generated poetry adheres to linguistic and stylistic constraints.\n",
    "\n",
    "## Key Components\n",
    "### **Language Model (LLM)**: \n",
    "At the core, an LLM is used to generate the actual lines of poetry based on a prompt. This model generates text while considering provided context and constraints.\n",
    "\n",
    "### **RAG (Retrieval Augmented Generation)**:\n",
    "This module wraps user prompts with context, such as examples of words that fit the required Qafiya (rhyme) and previous successful lines. It guides the LLM to generate poetry that follows the intended form.\n",
    "\n",
    "### **Extractor**:\n",
    "Automatically extracts Qafiya and Wazn from the generated line to validate them against the given constraints.\n",
    "\n",
    "### **Validators**:\n",
    "#### 1. **WaznValidator**:\n",
    "Ensures that the generated line conforms to the correct Wazn pattern.\n",
    "#### 2. **QafiyaValidator**:\n",
    "Ensures that the Qafiya (rhyme) of the generated line matches the expected one.\n",
    "\n",
    "### **Feedback Generator**:\n",
    "Provides feedback when the generated line doesn't produce the expected Qafiya or Wazn, which is sent back to the LLM to produce better lines.\n",
    "\n",
    "### **ShatrGenerator**:\n",
    "This is the central class responsible for orchestrating the poetry generation process. It interacts with the LLM, RAG, validators, and feedback generator to iteratively produce lines that meet the Wazn and Qafiya requirements.\n",
    "\n",
    "## **Workflow**:\n",
    "1. The user inputs a prompt.\n",
    "2. The program generates a line (called \"shatrs\") using the LLM.\n",
    "3. It extracts and validates the Wazn (rhythm) and Qafiya (rhyme) of the generated line.\n",
    "4. If the line is invalid, a feedback is generated, and we go back to step 2 with the new feedback.\n",
    "5. if the line is valid, we save the line and repeat the process until we create enough lines for a complete poem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Flow Diagram"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- ![Flow Diagram](flow-diagram.svg) -->\n",
    "<img src=\"flow-diagram.svg\" alt=\"Flow Diagram\" width=\"800\">\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Requirements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: ibm_cloud_sdk_core in ./.venv/lib/python3.10/site-packages (3.21.0)\n",
      "Requirement already satisfied: python-dateutil<3.0.0,>=2.8.2 in ./.venv/lib/python3.10/site-packages (from ibm_cloud_sdk_core) (2.9.0.post0)\n",
      "Requirement already satisfied: requests<3.0.0,>=2.31.0 in ./.venv/lib/python3.10/site-packages (from ibm_cloud_sdk_core) (2.32.2)\n",
      "Requirement already satisfied: urllib3<3.0.0,>=2.1.0 in ./.venv/lib/python3.10/site-packages (from ibm_cloud_sdk_core) (2.2.3)\n",
      "Requirement already satisfied: PyJWT<3.0.0,>=2.8.0 in ./.venv/lib/python3.10/site-packages (from ibm_cloud_sdk_core) (2.9.0)\n",
      "Requirement already satisfied: six>=1.5 in ./.venv/lib/python3.10/site-packages (from python-dateutil<3.0.0,>=2.8.2->ibm_cloud_sdk_core) (1.16.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in ./.venv/lib/python3.10/site-packages (from requests<3.0.0,>=2.31.0->ibm_cloud_sdk_core) (3.10)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in ./.venv/lib/python3.10/site-packages (from requests<3.0.0,>=2.31.0->ibm_cloud_sdk_core) (3.3.2)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in ./.venv/lib/python3.10/site-packages (from requests<3.0.0,>=2.31.0->ibm_cloud_sdk_core) (2024.8.30)\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install ibm_cloud_sdk_core"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LLM Interface\n",
    "An interface class for interacting with large language models. Any LLM implementation must define the `generate` method, which takes a prompt and generates text based on it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "from abc import ABC, abstractmethod\n",
    "class LLM_Interface(ABC):\n",
    "    @abstractmethod\n",
    "    def generate(self, prompt, **kwargs) -> str:\n",
    "        pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the `ALLAM` class which implements the `LLM_Interface` abstract class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "from ibm_cloud_sdk_core.authenticators import IAMAuthenticator\n",
    "\n",
    "BASE_URL = \"https://eu-de.ml.cloud.ibm.com/ml/\"\n",
    "\n",
    "class ALLAM(LLM_Interface):\n",
    "    def __init__(self, API_KEY):\n",
    "        self.model_id = \"sdaia/allam-1-13b-instruct\"\n",
    "        self.project_id = \"0a443bde-e9c6-41dc-b1f2-65c6292030e4\"\n",
    "\n",
    "        # get authentication token\n",
    "        authenticator = IAMAuthenticator(API_KEY)\n",
    "        token = authenticator.token_manager.get_token()\n",
    "        self.headers = {\n",
    "            'Accept': 'application/json',\n",
    "            'Content-Type': 'application/json',\n",
    "            'Authorization': f'Bearer {token}'\n",
    "        }\n",
    "\n",
    "        # set default parameters\n",
    "        self.parameters = {\n",
    "            \"decoding_method\": \"sample\",\n",
    "            \"max_new_tokens\": 30,\n",
    "            \"temperature\": 1,\n",
    "            \"top_k\": 50,\n",
    "            \"top_p\": 1,\n",
    "            \"repetition_penalty\": 2,\n",
    "        }\n",
    "\n",
    "    def generate(self, prompt, **kwargs):\n",
    "        url = BASE_URL + \"v1/text/generation?version=2024-08-30\"\n",
    "        self.body = {\n",
    "            \"input\": f\"<s> [INST] {prompt} [/INST]\",\n",
    "            \"model_id\": self.model_id,\n",
    "            \"project_id\": self.project_id,\n",
    "            \"parameters\": self.parameters\n",
    "        }\n",
    "\n",
    "        response = requests.post(url, headers=self.headers, json=self.body)\n",
    "        response.raise_for_status()\n",
    "        \n",
    "        data = response.json()\n",
    "        return data['results'][0]['generated_text']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RAG (Retrieval Augmented Generation)\n",
    "\n",
    "The RAG class provides context to the LLM by wrapping user input with example rhymes (Qafiya) and previous successful lines. It reads a database of words and processes them to suggest suitable rhyme words."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import random\n",
    "\n",
    "class RAG:\n",
    "    def __init__(self, filepath=\"qawafi-database.json\", qafiya=None):\n",
    "        self.qafiya = qafiya\n",
    "        self.message = None\n",
    "        self.db = None\n",
    "        with open(filepath, 'r', encoding=\"utf-8\") as f:\n",
    "            self.db = json.load(f)[\"data\"]\n",
    "        self.setQafiya(qafiya)\n",
    "    \n",
    "    def wrap(self, prompt, previous_shatrs=None, feedback=None):\n",
    "        full_text = \"\"\n",
    "        full_text += \"اكتب شطر واحد لجزء من قصيدة.\\n\"\n",
    "        \n",
    "        if self.qafiya:\n",
    "            full_text += f\" قافية القصيدة هي '{self.qafiya}'. \"\n",
    "\n",
    "        if self.message:\n",
    "            full_text += \"هنا بعض الامثلة لكلمات تنتهي بهذه القافية: \\n\"\n",
    "            # full_text += \"<QafiyaExamples>\"\n",
    "            full_text += \", \".join(random.sample(self.message, 10))\n",
    "            full_text += \"\\n\"\n",
    "            # full_text += \"</QafiyaExamples>\"\n",
    "        \n",
    "        if prompt:\n",
    "            full_text += f\"هنا الطلب اللي وضعه المستخدم:\\n\"\n",
    "            full_text += f\"{prompt}\\n\"\n",
    "        \n",
    "        if previous_shatrs:\n",
    "            full_text += \"هنا الشطور السابقة:\\n\"\n",
    "            full_text += f\"{previous_shatrs}\\n\"\n",
    "\n",
    "        if feedback:\n",
    "            full_text += \"هنا بعض النصائح على هذا اخر شطر تم ادخاله:\\n\"\n",
    "            full_text += f\"{feedback}\\n\"\n",
    "\n",
    "        return full_text\n",
    "    \n",
    "    def update(self, qafiya):\n",
    "        self.message = self.setQafiya(qafiya)\n",
    "        self.qafiya = qafiya\n",
    "\n",
    "    def setQafiya(self, qafiya):\n",
    "        if not qafiya:\n",
    "            return None\n",
    "\n",
    "        def processed(word):\n",
    "            if word.endswith(qafiya):\n",
    "                return word\n",
    "            if qafiya[-1] == \"ا\" and word.endswith(qafiya[0]):\n",
    "                return word+\"ا\"\n",
    "            if qafiya[-1] == \"ه\" and word.endswith(qafiya[0]):\n",
    "                return word+\"ه\"\n",
    "            if qafiya[-1] == \"ه\" and word.endswith(\"ة\"):\n",
    "                return word.replace(\"ة\" , f\"ه\")\n",
    "            return None\n",
    "        \n",
    "        output = []\n",
    "        for word in self.db:\n",
    "            new_word = processed(word)\n",
    "            if new_word:\n",
    "                output.append(new_word)\n",
    "\n",
    "        return output\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Demoing some `RAG` examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "اكتب شطر واحد لجزء من قصيدة.\n",
      "هنا الطلب اللي وضعه المستخدم:\n",
      "اكتب قصيدة عن امرؤ القيس\n",
      "\n",
      "اكتب شطر واحد لجزء من قصيدة.\n",
      " قافية القصيدة هي 'يا'. هنا بعض الامثلة لكلمات تنتهي بهذه القافية: \n",
      "عرايا, ينسيا, يغويا, يلفيا, مليشيا, قصايا, كوميديا, يبتغيا, شلايا, يدريا\n",
      "هنا الطلب اللي وضعه المستخدم:\n",
      "اكتب قصيدة عن امرؤ القيس\n",
      "\n",
      "اكتب شطر واحد لجزء من قصيدة.\n",
      " قافية القصيدة هي 'لا'. هنا بعض الامثلة لكلمات تنتهي بهذه القافية: \n",
      "صلولا, علولا, تكفلا, زعبيلا, يتخيلا, آجالا, حوافلا, زلا, خاذلا, كولسترولا\n",
      "هنا الطلب اللي وضعه المستخدم:\n",
      "اكتب قصيدة عن امرؤ القيس\n",
      "\n",
      "اكتب شطر واحد لجزء من قصيدة.\n",
      " قافية القصيدة هي 'لا'. هنا بعض الامثلة لكلمات تنتهي بهذه القافية: \n",
      "يوغلا, تواكلا, تأولا, مصاقلا, حثلا, تنصلا, يرتجلا, يزلزلا, يحنبلا, عبلا\n",
      "هنا الطلب اللي وضعه المستخدم:\n",
      "اكتب قصيدة عن امرؤ القيس\n",
      "هنا الشطور السابقة:\n",
      "['قَالَتْ فُطَيْمَةُ حَلِّ شِعْرَكَ مَدْحَهُ']\n",
      "\n",
      "اكتب شطر واحد لجزء من قصيدة.\n",
      " قافية القصيدة هي 'لا'. هنا بعض الامثلة لكلمات تنتهي بهذه القافية: \n",
      "علا, يجلا, اجتمالا, ندلا, احتفالا, يرهلا, يختبلا, سائلا, يتحنجلا, ثعلا\n",
      "هنا الطلب اللي وضعه المستخدم:\n",
      "اكتب قصيدة عن امرؤ القيس\n",
      "هنا الشطور السابقة:\n",
      "['قَالَتْ فُطَيْمَةُ حَلِّ شِعْرَكَ مَدْحَهُ', 'أَفَبَعْدَ كِنْدَةَ تَمْدَحَنَّ قَبِيلا']\n",
      "هنا بعض النصائح على هذا اخر شطر تم ادخاله:\n",
      "الشطر السابق لا يحتوي على قافية\n",
      "\n"
     ]
    }
   ],
   "source": [
    "rag = RAG()\n",
    "\n",
    "print(rag.wrap(\"اكتب قصيدة عن امرؤ القيس\"))\n",
    "rag.update(\"يا\")\n",
    "print(rag.wrap(\"اكتب قصيدة عن امرؤ القيس\"))\n",
    "rag.update(\"لا\")\n",
    "print(rag.wrap(\"اكتب قصيدة عن امرؤ القيس\"))\n",
    "print(rag.wrap(\"اكتب قصيدة عن امرؤ القيس\", [\"قَالَتْ فُطَيْمَةُ حَلِّ شِعْرَكَ مَدْحَهُ\"]))\n",
    "print(rag.wrap(\"اكتب قصيدة عن امرؤ القيس\", [\"قَالَتْ فُطَيْمَةُ حَلِّ شِعْرَكَ مَدْحَهُ\", \"أَفَبَعْدَ كِنْدَةَ تَمْدَحَنَّ قَبِيلا\"], \"الشطر السابق لا يحتوي على قافية\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extractor\n",
    "The Extractor class is responsible for extracting both the Qafiya and Wazn types from a generated line (shatr). This information is then used by the validators to ensure the correctness of the Qafiya and Wazn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Extractor:\n",
    "    def __init__(self):\n",
    "        pass\n",
    "\n",
    "    def extract(self, shatr):\n",
    "        qafiya_type = self.extract_qafiya(shatr)\n",
    "        wazn_type = self.extract_wazn(shatr)\n",
    "        return qafiya_type, wazn_type\n",
    "\n",
    "    def extract_qafiya(self, shatr):\n",
    "        # TODO\n",
    "        return None\n",
    "\n",
    "    def extract_wazn(self, shatr):\n",
    "        # TODO\n",
    "        return None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Demoing some `Extractor` examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(None, None)\n",
      "(None, None)\n"
     ]
    }
   ],
   "source": [
    "extractor = Extractor()\n",
    "\n",
    "print(extractor.extract(\"قَالَتْ فُطَيْمَةُ حَلِّ شِعْرَكَ مَدْحَهُ\"))\n",
    "print(extractor.extract(\"أَفَبَعْدَ كِنْدَةَ تَمْدَحَنَّ قَبِيلا\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## WaznValidator\n",
    "This class is responsible for validating the Wazn of a given shatr. It compares the Wazn of the current shatr with the previous one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WaznValidator:\n",
    "    def __init__(self):\n",
    "        pass\n",
    "\n",
    "    def validate_wazn(self, current_wazn, previous_wazn=None):\n",
    "        # TODO: implement validation\n",
    "        # if previous_wazn is None:\n",
    "        #     return True\n",
    "        # if current_wazn == previous_wazn:\n",
    "        #     return True\n",
    "        # else:\n",
    "        #     return False\n",
    "        return True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Demoing some `WaznValidator` examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "wazn_validator = WaznValidator()\n",
    "\n",
    "print(wazn_validator.validate_wazn(\"طويل\"))\n",
    "print(wazn_validator.validate_wazn(\"طويل\", \"بسيط\"))\n",
    "print(wazn_validator.validate_wazn(\"طويل\", \"طويل\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### QafiyaValidator\n",
    "The QafiyaValidator class checks if the generated Qafiya matches the expected one. This ensures that the lines end with the correct sounds as required by the given Qafiya."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [],
   "source": [
    "class QafiyaValidator:\n",
    "    def __init__(self):\n",
    "        pass\n",
    "\n",
    "    def validate_qafiya(self, current_qafiya, previous_qafiya=None):\n",
    "        # TODO: implement validation\n",
    "        # if previous_qafiya is None:\n",
    "        #     return True        \n",
    "        # if current_qafiya == previous_qafiya:\n",
    "        #     return True\n",
    "        # else:\n",
    "        #     return False\n",
    "        return True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Demoing some `QafiyaValidator` examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "qafiya_validator = QafiyaValidator()\n",
    "\n",
    "print(qafiya_validator.validate_qafiya(\"ه\"))\n",
    "print(qafiya_validator.validate_qafiya(\"ل\", \"ا\"))\n",
    "print(qafiya_validator.validate_qafiya(\"ل\", \"ل\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FeedbackGenerator\n",
    "\n",
    "This class generates feedback for incorrect or invalid Wazn or Qafiya. It provides suggestions based on what went wrong to guide further iterations of line generation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedbackGenerator:\n",
    "    def generate_feedback(self, type, invalid_item, expected_item, invalid_shatr,):\n",
    "        if type == \"qafiya\":\n",
    "            return self.qafiya_feedback(invalid_item, expected_item, invalid_shatr)\n",
    "        \n",
    "        if type == \"wazn\":\n",
    "            return self.wazn_feedback(invalid_item, expected_item, invalid_shatr)\n",
    "    \n",
    "    def qafiya_feedback(self, invalid_qafiya, expected_qafiya, invalid_shatr):\n",
    "        # TODO\n",
    "\n",
    "        text = \"\"\n",
    "        text += f\"الشطر المدخل {invalid_shatr} غير صحيح. \"\n",
    "        text += f\"القافية المدخلة {invalid_qafiya} غير صحيحة. \"\n",
    "        text += f\"القافية الصحيحة هي {expected_qafiya}. \"\n",
    "        \n",
    "        return text\n",
    "\n",
    "    def wazn_feedback(self, invalid_wazn, expected_wazn, invalid_shatr):\n",
    "        # TODO\n",
    "\n",
    "        text = \"\"\n",
    "        text += f\"الشطر المدخل {invalid_shatr} غير صحيح. \"\n",
    "        text += f\"الوزن المدخل {invalid_wazn} غير صحيح. \"\n",
    "        text += f\"الوزن الصحيح هو {expected_wazn}. \"\n",
    "\n",
    "        return text"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Demoing some `FeedbackGenerator` examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "القافية المدخلة ق غير صحيحة. القافية الصحيحة هي ل. \n",
      "الوزن المدخل بسيط غير صحيح. الوزن الصحيح هو كامل. \n"
     ]
    }
   ],
   "source": [
    "generator = FeedbackGenerator()\n",
    "\n",
    "print(generator.generate_feedback(\"qafiya\", \"ق\", \"ل\", \"أَفَبَعْدَ كِنْدَةَ تَمْدَحَنَّ قَبِيلا\"))\n",
    "print(generator.generate_feedback(\"wazn\", \"بسيط\", \"كامل\", \"قَالَتْ فُطَيْمَةُ حَلِّ شِعرَكَ مَدحَهُ\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ShatrGenerator\n",
    "\n",
    "The `ShatrGenerator` class generates lines of poetry based on the given prompt. It uses a language model (LLM), a feedback mechanism, and validators for Wazn and Qafiya. The process involves:\n",
    "\n",
    "- Generating a line (shatr) from the LLM class, which implements `LLM_Interface`.\n",
    "- Extracting the Wazn and Qafiya from the `Extractor`\n",
    "- Validating both the Wazn and Qafiya using `WaznValidator` and `QafiyaValidator`, respectively.\n",
    "- Give feedback to the LLM for incorrect Wazn or Qafiya using `FeedbackGenerator`.\n",
    "- Keep trying until it generates a correct shatr."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ShatrGenerator:\n",
    "    def __init__(self, llm, rag=None, feedback_generator=None, extractor=None, wazn_validator=None, qafiya_validator=None):\n",
    "        self.llm = llm\n",
    "        self.rag = rag or RAG()\n",
    "        self.feedback_generator = feedback_generator or FeedbackGenerator()\n",
    "        self.extractor = extractor or Extractor()\n",
    "        self.wazn_validator = wazn_validator or WaznValidator()\n",
    "        self.qafiya_validator = qafiya_validator or QafiyaValidator()\n",
    "        \n",
    "    def generate_shatr(self, prompt, wazn=None, qafiya=None, feedback=None, previous_shatrs=None):\n",
    "        valid = False\n",
    "        \n",
    "        while not valid:\n",
    "            # Step 1: Generate a shatr\n",
    "            shatr = self.llm.generate(self.rag.wrap(prompt, previous_shatrs, feedback))\n",
    "            print(shatr)\n",
    "\n",
    "            # Step 2: Extract Wazn and Qafiya\n",
    "            new_qafiya, new_wazn = self.extractor.extract(shatr)\n",
    "\n",
    "            # Step 3: Validate Wazn\n",
    "            valid_wazn = self.wazn_validator.validate_wazn(new_wazn, wazn)\n",
    "            if not valid_wazn:\n",
    "                feedback = self.feedback_generator.generate_feedback(\"wazn\", shatr, wazn=new_wazn)\n",
    "                continue  # Loop back to regenerate\n",
    "\n",
    "            wazn = new_wazn\n",
    "\n",
    "            # Step 5: Validate Qafiya\n",
    "            valid_qafiya = self.qafiya_validator.validate_qafiya(new_qafiya, qafiya)\n",
    "            if not valid_qafiya:\n",
    "                feedback = self.feedback_generator.generate_feedback(\"qafiya\", shatr, qafiya=new_qafiya)\n",
    "                continue  # Loop back to regenerate\n",
    "\n",
    "            # Step 6: Update RAG and finalize shatr\n",
    "            self.rag.update(qafiya)\n",
    "            valid = True\n",
    "        return shatr\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### generate_qasida (function)\n",
    "\n",
    "This function generates a full qasida (poem) by producing individual lines of poetry using the `ShatrGenerator`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "def infer_wazn(prompt):\n",
    "    return None\n",
    "\n",
    "def infer_qafiya(prompt):\n",
    "    return None\n",
    "\n",
    "def infer_length(prompt):\n",
    "    return 6\n",
    "\n",
    "def generate_qasida(prompt, shatr_generator):\n",
    "    wazn = infer_wazn(prompt)\n",
    "    qafiya = infer_qafiya(prompt)\n",
    "    length = infer_length(prompt)\n",
    "    \n",
    "    shatrs = []\n",
    "    for shatr_idx in range(length):\n",
    "        shatr = shatr_generator.generate_shatr(prompt, wazn, qafiya, shatrs)\n",
    "        shatrs.append(shatr)\n",
    "    \n",
    "    # output = \"\"\n",
    "    # for i, shdr in enumerate(shatrs):\n",
    "    #     output += shdr + (\"\\n\" if i % 2 else \" # \")\n",
    "    \n",
    "    return shatrs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demonstration\n",
    "\n",
    "Running the code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is a fake LLM class to test the program\n",
    "\n",
    "class FakeLLM(LLM_Interface):\n",
    "    def __init__(self):\n",
    "        self.i = 0\n",
    "        self.lines = '''قَالَتْ فُطَيْمَةُ حَلِّ شِعْرَكَ مَدْحَهُ\n",
    "أَفَبَعْدَ كِنْدَةَ تَمْدَحَنَّ قَبِيلا\n",
    "وَهَمُ الكِرَامُ بَنُو الخَضَارِمَةِ العُلا\n",
    "ِسَمَيْدَعٍ أَكْرِمْ بِذَاكَ نَجِيلا\n",
    "يَا أَيُّها السَّاعِي لِيُدْرِكَ مَجْدَنَا\n",
    "ثَكِلَتْكَ أُمُّكَ هَلْ تَرُدُّ قَتِيلا\n",
    "هَلْ تَرْقَيَنَّ إلى السَّماءِ بِسُلَّمٍ\n",
    "وَلَتَرْجِعَنَّ إلى العَزِيزِ ذَلِيلا\n",
    "سَائِلْ بَنِي مَلِكِ المُلُوكِ إذا الْتَقَوا\n",
    "َنَّا وَعَنْكُمْ لا تَعَاشَ جَهُولا\n",
    "مِنَّا الذي مَلِكَ المَعَاشِرَ عَنْوَةً\n",
    "مَلَكَ الفَضَاءَ فَسَلْ بِذَاك عُقُولا\n",
    "وَبَنُوهُ قَدْ مَلَكُوا خِلافَةَ مُلْكِهِ\n",
    "شُبَّانَ حَرْبٍ سَادَةً وَكُهُولا\n",
    "قالوا لَهُ : هَلْ أنتَ قَاضٍ ما تَرَى\n",
    "إِنَّا نَرَى لَكَ ذا المَقَامَ قَلِيلا\n",
    "فَقَضَى لكلِّ قَبِيلةٍ بِتِرَاتِهِمْ\n",
    "لَمْ يَأْلُهُمْ في مُلْكِهمْ تَعْدِيلا\n",
    "فَثَوَى وَوَرَّثَ مُلْكِ مَنْ وَطِئَ الحَصَى\n",
    "قَسْرًا أبوهُ عَنْوَةً وَنُحُولا\n",
    "سَائِلْ بَنِي أَسَدٍ بِمَقْتَلِ رَبِّهِمْ\n",
    "حُجْرِ بنِ أُمِّ قَطَامِ جَلَّ قَتِيلا\n",
    "إذا سَارَ ذو التَّاجِ الهِجَانِ بِجَحْفَلٍ\n",
    "لَجِبٍ يُجَاوَبُ بالفَلاةِ صَهِيلا\n",
    "حتى أَبَالَ الخَيْلَ في عَرَصَاتِهِمْ\n",
    "فَشَفَى وَزَادَ على الشِّفَاءِ غَلِيلا\n",
    "أَحْمَى دُرُوعَهُمُ فَسَرْبَلَهُمْ بِهَا\n",
    "والنَّارَ كَحَّلَهُمْ بها تَكْحِيلا\n",
    "وأقامَ يَسْقِي الرَّاحِ في هَامَاتِهِمْ    \n",
    "مَلِكٌ يُعَلُّ بِشُرْبِها تَعْلِيلا\n",
    "والبِيْضَ قَنَّعَهَا شَدِيدًا حَرُّهُا\n",
    "فَكَفَى بذلكَ لِلْعِدَا تَنْكِيلا\n",
    "حَلَّتْ لَهُ مِنْ بَعْدِ تَحْرِيمٍ لَهَا\n",
    "أَو أَنْ يَمَسَّ الرَّأسَ منه غُسُولا\n",
    "حتى أباحَ ديارَهمْ فَأَبَارَهُمْ\n",
    "فَعَمُوا فهمْ لا يَهْتَدونَ سَبِيلا'''.split(\"\\n\")\n",
    "    \n",
    "\n",
    "    def generate(self, prompt, **kwargs):\n",
    "        line = self.lines[self.i].strip()\n",
    "        self.i = (self.i + 1) % len(self.lines)\n",
    "        return line\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "قَالَتْ فُطَيْمَةُ حَلِّ شِعْرَكَ مَدْحَهُ\n",
      "أَفَبَعْدَ كِنْدَةَ تَمْدَحَنَّ قَبِيلا\n",
      "وَهَمُ الكِرَامُ بَنُو الخَضَارِمَةِ العُلا\n",
      "ِسَمَيْدَعٍ أَكْرِمْ بِذَاكَ نَجِيلا\n",
      "يَا أَيُّها السَّاعِي لِيُدْرِكَ مَجْدَنَا\n",
      "ثَكِلَتْكَ أُمُّكَ هَلْ تَرُدُّ قَتِيلا\n",
      "قَالَتْ فُطَيْمَةُ حَلِّ شِعْرَكَ مَدْحَهُ\n",
      "أَفَبَعْدَ كِنْدَةَ تَمْدَحَنَّ قَبِيلا\n",
      "وَهَمُ الكِرَامُ بَنُو الخَضَارِمَةِ العُلا\n",
      "ِسَمَيْدَعٍ أَكْرِمْ بِذَاكَ نَجِيلا\n",
      "يَا أَيُّها السَّاعِي لِيُدْرِكَ مَجْدَنَا\n",
      "ثَكِلَتْكَ أُمُّكَ هَلْ تَرُدُّ قَتِيلا\n"
     ]
    }
   ],
   "source": [
    "# api_key = input(\"Enter API key: \")\n",
    "# llm = ALLAM(api_key)\n",
    "llm = FakeLLM()\n",
    "shatr_generator = ShatrGenerator(llm)\n",
    "while True:\n",
    "    prompt = input(\"Enter a prompt (type 'exit' to stop): \")\n",
    "    if not prompt or prompt == \"exit\":\n",
    "        break\n",
    "    qasida = generate_qasida(prompt, shatr_generator)\n",
    "    print(\"\\n\".join(qasida))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Viewing the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "قَالَتْ فُطَيْمَةُ حَلِّ شِعْرَكَ مَدْحَهُ # أَفَبَعْدَ كِنْدَةَ تَمْدَحَنَّ قَبِيلا\n",
      " وَهَمُ الكِرَامُ بَنُو الخَضَارِمَةِ العُلا # ِسَمَيْدَعٍ أَكْرِمْ بِذَاكَ نَجِيلا\n",
      " يَا أَيُّها السَّاعِي لِيُدْرِكَ مَجْدَنَا # ثَكِلَتْكَ أُمُّكَ هَلْ تَرُدُّ قَتِيلا\n",
      " "
     ]
    }
   ],
   "source": [
    "for i, line in enumerate(qasida):\n",
    "    text = line + (\"\\n\" if i % 2 else \" #\")\n",
    "    print(text, end=\" \")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}