Update train.py

train.py

@@ -1,12 +1,5 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-import os
-
-# Avoid permission errors in Hugging Face Spaces
-os.environ["HF_HOME"] = "/workspace/.cache"
-os.environ["HF_DATASETS_CACHE"] = "/workspace/.cache/hf_datasets"
-os.environ["TRANSFORMERS_CACHE"] = "/workspace/.cache/hf_transformers"
-
 from datasets import load_dataset
 from trl import GRPOConfig, GRPOTrainer
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -51,202 +44,4 @@ roman_to_int = {'I': 1, 'V': 5, 'X': 10, 'L': 50, 'C': 100, 'D': 500, 'M': 1000}
 print("Connecting to ChromaDB...")
 chroma_client = HttpClient(host="34.126.200.250", port=8000)
 collection = chroma_client.get_collection(name="orwell_books")
-nlp = spacy.load("en_core_web_sm")
-phrase_matcher = PhraseMatcher(nlp.vocab, attr="LOWER")
-patterns = [nlp.make_doc(title) for title in BOOK_TITLES]
-phrase_matcher.add("BOOK_TITLE", patterns)
 print("Connected to ChromaDB.")
-
-def extract_metadata(prompt):
-    metadata = {"$and": [{"title": None}]}
-    doc = nlp(prompt)
-    part_no = get_part_number(prompt)
-    name = None
-    matches = phrase_matcher(doc)
-    for _, start, end in matches:
-        span = doc[start:end]
-        metadata["$and"][0]["title"] = span.text.lower()
-        name = span.text.lower()
-        break
-    if part_no is not None:
-        digit_word = digit_to_spelled[part_no]
-        roman_word = digit_to_roman[part_no]
-        chapter_filter = {
-            "$or": [
-                {"chapter": {"$in": [f"chapter {part_no}", f"part {part_no}", f"chapter {digit_word}", f"part {digit_word}", f"chapter {roman_word}"]}},
-                {"subchapter": {"$in": [f"chapter {part_no}", f"part {part_no}", f"chapter {digit_word}", f"part {digit_word}", f"chapter {roman_word}"]}}
-            ]
-        }
-        metadata["$and"].append(chapter_filter)
-    else:
-        metadata = {"title": name}
-    return metadata, name
-
-def retrieve_context(prompt, top_k=1):
-    filters, nameBook = extract_metadata(prompt)
-    results = collection.query(
-        query_texts=["Give me the book named " + nameBook],
-        where=filters,
-        n_results=top_k
-    )
-    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])
-    ]
-
-def detect_question_type(user_prompt):
-    prompt_lower = user_prompt.lower()
-    if any(k in prompt_lower for k in ["mcq", "multiple choice"]): return "MCQs"
-    if any(k in prompt_lower for k in ["essay", "long answer"]): return "Essay"
-    if any(k in prompt_lower for k in ["short answer"]): return "Short Answer"
-    if any(k in prompt_lower for k in ["paragraph"]): return "Paragraph Answer"
-    if any(k in prompt_lower for k in ["passage"]): return "Passage-Based"
-    if "quiz" in prompt_lower: return "Quiz"
-    return "Short Answer"
-
-def roman_to_number(roman):
-    total, prev_value = 0, 0
-    for char in reversed(roman.upper()):
-        value = roman_to_int.get(char, 0)
-        total += -value if value < prev_value else value
-        prev_value = value
-    return total if total > 0 else None
-
-def get_part_number(prompt):
-    doc = nlp(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):
-            for j in range(1, 3):
-                if i + j < len(doc) and doc[i + j].text.lower() in part_keywords:
-                    return spelled_to_digit.get(token_lower) or ordinal_to_digit.get(token_lower) or int(token_lower) if number_pattern.match(token_lower) else roman_to_number(token_lower)
-        elif token_lower in part_keywords:
-            for j in range(1, 4):
-                if i + j < len(doc):
-                    next_token = doc[i + j].text.lower()
-                    return spelled_to_digit.get(next_token) or ordinal_to_digit.get(next_token) or int(next_token) if number_pattern.match(next_token) else roman_to_number(next_token)
-    return None
-
-def get_prompt(contexts, user_prompt=""):
-    question_type = detect_question_type(user_prompt)
-    citation_instr = """
-      - Every question and explanation must include at least one **inline citation** in the format: [1]
-      - 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_block}
-      Supporting context
-
-      Sources:
-      [1] Book Title, Chapter X
-    """
-    prompt = f"""
-      You are a literature teaching assistant. Generate one {question_type} question based on the context.
-      {citation_instr}
-
-      Context:
-      {contexts[0]['content']}
-
-      Now Complete:
-      {output_format}
-    """
-    return prompt
-
-def make_format(example):
-    return {
-        "prompt": [
-            {"role": "system", "content": get_prompt(retrieve_context(example["prompt"]), example["prompt"])} ,
-            {"role": "user", "content": example["prompt"]},
-        ],
-    }
-
-def binary_reward(value): return 1.0 if value == "yes" else -1.0
-
-def normalize_rating(value, max_value=5): return value / max_value
-
-def compute_reward(sample):
-    intent = binary_reward(sample.get("intent_relevance.responses", ["no"])[0])
-    citation = binary_reward(sample.get("citation_support.responses", ["no"])[0])
-    hallucination = binary_reward(sample.get("hallucination_check.responses", ["no"])[0])
-    clarity = normalize_rating(sample.get("clarity_rating.responses", [0])[0])
-    relevance = normalize_rating(sample.get("relevance_rating.responses", [0])[0])
-    overall = normalize_rating(sample.get("overall_quality.responses", [0])[0])
-    return 0.2*intent + 0.2*citation + 0.2*hallucination + 0.1*clarity + 0.15*relevance + 0.15*overall
-
-def add_reward(example):
-    example["reward"] = compute_reward(example)
-    return example
-
-def reward_calculate(completions, **kwargs):
-    return kwargs["reward"]
-
-print("Loading raw dataset...")
-raw_dataset = load_dataset(DATA_ID, split="train")
-train_dataset = raw_dataset.map(make_format).map(add_reward).remove_columns([
-    'id', 'status', 'inserted_at', 'updated_at', '_server_id',
-    'intent_relevance.responses.users', 'intent_relevance.responses.status',
-    'citation_support.responses.users', 'citation_support.responses.status',
-    'hallucination_check.responses.users', 'hallucination_check.responses.status',
-    'clarity_rating.responses.users', 'clarity_rating.responses.status',
-    'relevance_rating.responses.users', 'relevance_rating.responses.status',
-    'overall_quality.responses.users', 'overall_quality.responses.status'])
-
-print("Loading model and tokenizer...")
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto", device_map="auto")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-tokenizer.save_pretrained(MODEL_ID + "_test")
-
-print("Applying LoRA...")
-lora_config = LoraConfig(r=32, lora_alpha=64, target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], lora_dropout=0.05, bias="none", task_type="CAUSAL_LM")
-model = get_peft_model(model, lora_config)
-model.print_trainable_parameters()
-
-print("Setting up GRPO training...")
-training_args = GRPOConfig(
-    output_dir=MODEL_ID + "_test",
-    learning_rate=2e-5,
-    remove_unused_columns=False,
-    gradient_accumulation_steps=16,
-    num_train_epochs=3,
-    bf16=True,
-    max_completion_length=96,
-    num_generations=4,
-    max_prompt_length=256,
-    report_to=["tensorboard"],
-    logging_steps=10,
-    push_to_hub=True,
-    save_strategy="steps",
-    save_steps=50,
-)
-
-trainer = GRPOTrainer(
-    model=model,
-    reward_funcs=reward_calculate,
-    args=training_args,
-    train_dataset=train_dataset
-)
-
-print("Starting training...")
-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!")