RebeccaL25/EventExtraction

1from fastapi import FastAPI2from fastapi.middleware.cors import CORSMiddleware3from pydantic import BaseModel4import spacy5from spacy.tokens import Span6import dateparser7import re8from datetime import timedelta9 10# Load trained spaCy model11nlp = spacy.load("output/model-last")12 13app = FastAPI()14 15# Enable CORS16app.add_middleware(17    CORSMiddleware,18    allow_origins=["*"],19    allow_credentials=True,20    allow_methods=["*"],21    allow_headers=["*"],22)23 24class TextInput(BaseModel):25    text: str26 27# Constants28MONTHS = {"jan", "feb", "mar", "apr", "may", "jun", "jul", "aug", "sep", "oct", "nov", "dec"}29DATE_REGEX = re.compile(r"\b\d{1,2}[/-]\d{1,2}([/-]\d{2,4})?\b", re.IGNORECASE)30DURATION_KEYWORDS = {"hour", "hours", "minutes", "mins", "min"}31 32def normalize_range(text):33    if not text:34        return ""35    text = text.lower().strip()36    text = re.sub(r'\b(from|between)\b', '', text)37    text = re.sub(r'\s*\.\.\.\s*', ' - ', text)38    text = re.sub(r'\s*(?:to|until|through|and|–|-)\s*', ' - ', text)39    return text40 41def fix_time_format(time_str):42    return re.sub(43        r'^(\d{3,4})(am|pm)?$',44        lambda m: f"{m.group(1)[:-2]}:{m.group(1)[-2:]} {m.group(2) or ''}".strip(),45        time_str46    )47 48def parse_duration_phrase(text):49    text = text.lower().strip()50 51    text = re.sub(r'(\d+(?:\.\d+)?(?:hr|hrs|hour|hours))(?=\d)', r'\1 ', text)52 53    match = re.search(54        r'(?P<start>[\d:apm\s]+?)\s*(?:for|in)?\s*'55        r'(?P<value1>\d+\.?\d*)\s*(?P<unit1>hr|hrs|hour|hours|min|mins|minutes)\b'56        r'(?:\s*(?P<value2>\d+\.?\d*)\s*(?P<unit2>hr|hrs|hour|hours|min|mins|minutes)\b)?',57        text58    )59 60    if not match:61        return "", "", False62 63    start_raw = match.group("start").strip()64    duration1 = float(match.group("value1"))65    unit1 = match.group("unit1")66 67    duration2 = float(match.group("value2")) if match.group("value2") else 068    unit2 = match.group("unit2") if match.group("unit2") else ""69 70    start_time = dateparser.parse(fix_time_format(start_raw))71    if not start_time:72        return "", "", False73 74    # Calculate total duration75    total_duration = timedelta()76    if unit1 in {"hour", "hours", "hr", "hrs"}:77        total_duration += timedelta(hours=duration1)78    else:79        total_duration += timedelta(minutes=duration1)80 81    if unit2:82        if unit2 in {"hour", "hours", "hr", "hrs"}:83            total_duration += timedelta(hours=duration2)84        else:85            total_duration += timedelta(minutes=duration2)86 87    end_time = start_time + total_duration88 89    return start_time.strftime("%H:%M"), end_time.strftime("%H:%M"), True90 91def parse_date_range(text):92    normalized = normalize_range(text)93    parts = [p.strip() for p in normalized.split(' - ') if p.strip()]94    parsed = []95    for part in parts:96        dt = dateparser.parse(part)97        if dt:98            parsed.append(dt.strftime('%Y-%m-%d'))99    return parsed if parsed else ["", ""]100 101def parse_time_range(text):102    # Try detecting "start for duration" pattern first103    start, end, matched = parse_duration_phrase(text)104    if matched:105        return [start, end]106 107    normalized = normalize_range(text)108    parts = [p.strip() for p in normalized.split(' - ') if p.strip()]109    parsed = []110    for part in parts:111        fixed = fix_time_format(part)112        dt = dateparser.parse(fixed)113        if dt:114            parsed.append(dt.strftime('%H:%M'))115    return parsed if parsed else ["", ""]116 117def correct_ner_labels(doc):118    corrected_ents = []119    for ent in doc.ents:120        text = ent.text.lower()121        label = ent.label_122 123        if label == "time_range":124            if DATE_REGEX.search(text) or any(m in text for m in MONTHS):125                corrected_ents.append(Span(doc, ent.start, ent.end, label="date_range"))126                continue127 128        if label == "date_range":129            if any(keyword in text for keyword in DURATION_KEYWORDS):130                corrected_ents.append(Span(doc, ent.start, ent.end, label="time_range"))131                continue132 133        corrected_ents.append(ent)134    doc.ents = corrected_ents135    return doc136 137@app.post("/extract")138def extract_entities(input: TextInput):139    doc = nlp(input.text)140    doc = correct_ner_labels(doc)141 142    result = {143        "event_name": "",144        "date_range": "",145        "time_range": "",146        "loc": "",147        "freq": "",148        "parsed_start_date": "",149        "parsed_end_date": "",150        "parsed_start_time": "",151        "parsed_end_time": "",152        "recurrence_pattern": ""153    }154 155    for ent in doc.ents:156        label = ent.label_157        text = ent.text.strip()158        if label in result:159            result[label] = text160 161    # Parse dates162    if result["date_range"]:163        dates = parse_date_range(result["date_range"])164        if len(dates) >= 1:165            result["parsed_start_date"] = dates[0]166        if len(dates) >= 2:167            result["parsed_end_date"] = dates[1]168 169    # Parse times including durations170    if result["time_range"]:171        times = parse_time_range(result["time_range"])172        if len(times) >= 1:173            result["parsed_start_time"] = times[0]174        if len(times) >= 2:175            result["parsed_end_time"] = times[1]176 177    # Frequency parsing → recurrence pattern178    if result["freq"]:179        freq_lower = result["freq"].lower()180        if "every day" in freq_lower or "daily" in freq_lower:181            result["recurrence_pattern"] = "daily"182        elif "weekend" in freq_lower:183            result["recurrence_pattern"] = "weekend"184        elif "week day" in freq_lower or "weekday" in freq_lower:185            result["recurrence_pattern"] = "weekday"186        else:187            for day in ["monday", "tuesday", "wednesday", "thursday", "friday", "saturday", "sunday"]:188                if day in freq_lower:189                    result["recurrence_pattern"] = day190                    break191 192    return result193 194@app.get("/")195def root():196    return {"status": "EventExtraction FastAPI is running"}197