| |
| |
| from datasets import load_dataset |
| from trl import GRPOConfig, GRPOTrainer |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| from peft import LoraConfig, get_peft_model |
| from chromadb import HttpClient |
| import torch |
| import spacy |
| from spacy.matcher import PhraseMatcher |
| import re |
|
|
| |
| MODEL_ID = "Qwen/Qwen2.5-7B-Instruct" |
| DATA_ID = "tapxc3/owast_new" |
| BOOK_TITLES = [ |
| "Nineteen Eighty-Four", "Animal Farm", "Shooting an Elephant", "Homage to Catalonia", |
| "Coming Up for Air", "Down and Out in Paris and London", "Keep the Aspidistra Flying", |
| "A Clergyman's Daughter", "Fifty Orwell Essays", "Politics and the English Language", |
| "The Road to Wigan Pier", "Burmese Days" |
| ] |
|
|
| 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, |
| "eighteen": 18, "nineteen": 19, "twenty": 20} |
|
|
| digit_to_spelled = {v: k for k, v in spelled_to_digit.items()} |
|
|
| 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, |
| "seventh": 7, "eighth": 8, "ninth": 9, "tenth": 10, "eleventh": 11, "twelfth": 12, |
| "thirteenth": 13, "fourteenth": 14, "fifteenth": 15, "sixteenth": 16, "seventeenth": 17, |
| "eighteenth": 18, "nineteenth": 19, "twentieth": 20 |
| } |
|
|
| 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") |
| print("Connected to ChromaDB.") |
|
|
