Update train.py

train.py

@@ -21,7 +21,6 @@ BOOK_TITLES = [
     "The Road to Wigan Pier", "Burmese Days"
 ]
 
-# Maps for conversion
 spelled_to_digit = {"one": 1, "two": 2, "three": 3, "four": 4, "five": 5, "six": 6,
                    "seven": 7, "eight": 8, "nine": 9, "ten": 10, "eleven": 11, "twelve": 12,
                    "thirteen": 13, "fourteen": 14, "fifteen": 15, "sixteen": 16, "seventeen": 17,
@@ -29,11 +28,10 @@ spelled_to_digit = {"one": 1, "two": 2, "three": 3, "four": 4, "five": 5, "six":
 
 digit_to_spelled = {v: k for k, v in spelled_to_digit.items()}
 
-digit_to_roman = {
-    1: "i", 2: "ii", 3: "iii", 4: "iv", 5: "v", 6: "vi", 7: "vii", 8: "viii", 9: "ix", 10: "x",
-    11: "xi", 12: "xii", 13: "xiii", 14: "xiv", 15: "xv", 16: "xvi", 17: "xvii", 18: "xviii",
-    19: "xix", 20: "xx"
-}
+digit_to_roman = {v: r for v, r in enumerate(
+    ["i", "ii", "iii", "iv", "v", "vi", "vii", "viii", "ix", "x", "xi", "xii", "xiii", "xiv", "xv", "xvi", "xvii", "xviii", "xix", "xx"],
+    start=1
+)}
 
 ordinal_to_digit = {
     "first": 1, "second": 2, "third": 3, "fourth": 4, "fifth": 5, "sixth": 6,
@@ -44,7 +42,6 @@ ordinal_to_digit = {
 
 roman_to_int = {'I': 1, 'V': 5, 'X': 10, 'L': 50, 'C': 100, 'D': 500, 'M': 1000}
 
-# ChromaDB setup
 print("Connecting to ChromaDB...")
 chroma_client = HttpClient(host="34.126.200.250", port=8000)
 collection = chroma_client.get_collection(name="orwell_books")
@@ -88,7 +85,6 @@ def retrieve_context(prompt, top_k=1):
     )
     if not results['documents'] or not results['documents'][0]:
         return [{"content": "No relevant context found.", "meta": {}}]
-
     return [
         {"content": doc, "meta": meta}
         for doc, meta in zip(results['documents'][0], results['metadatas'][0])
@@ -117,7 +113,6 @@ def get_part_number(prompt):
     part_keywords = {"chapter", "part", "section"}
     number_pattern = re.compile(r"^\\d+$")
     roman_pattern = re.compile(r"^[IVXLCDM]+$", re.IGNORECASE)
-
     for i, token in enumerate(doc):
         token_lower = token.text.lower()
         if token_lower in spelled_to_digit or token_lower in ordinal_to_digit or roman_pattern.match(token_lower) or number_pattern.match(token_lower):
@@ -138,9 +133,19 @@ def get_prompt(contexts, user_prompt=""):
       - Each inline citation must correspond to a **source reference at the end** of the response, using this exact format:
         [1] Book Title, Chapter X
     """
+    if question_type == "MCQs":
+        answer_block = (
+            "Answer options:\n"
+            "  A. Option A\n"
+            "  B. Option B\n"
+            "  C. Option C\n"
+            "  D. Option D\n"
+        )
+    else:
+        answer_block = ""
     output_format = f"""
       Question
-      {'Answer options:\n  A. Option A\n  B. Option B\n  C. Option C\n  D. Option D\n' if question_type == 'MCQs' else ''}
+      {answer_block}
       Supporting context
 
       Sources:
@@ -186,7 +191,6 @@ def add_reward(example):
 def reward_calculate(completions, **kwargs):
     return kwargs["reward"]
 
-# Load Dataset
 print("Loading raw dataset...")
 raw_dataset = load_dataset(DATA_ID, split="train")
 train_dataset = raw_dataset.map(make_format).map(add_reward).remove_columns([
@@ -239,4 +243,4 @@ trainer.train()
 print("Saving and pushing model...")
 trainer.save_model(training_args.output_dir)
 trainer.push_to_hub(dataset_name=DATA_ID)
-print("Training complete and model uploaded!")
+print("Training complete and model uploaded!")