import gradio as gr

from transformers import AutoTokenizer,AutoModelForCausalLM
import torch
import re

# Load base LLaMA2 model and fine-tuned adapter

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("NousResearch/Llama-2-7b-chat-hf" )

# Load base model (FP16 + auto device mapping for GPU support)
base_model = AutoModelForCausalLM.from_pretrained(
        "Zoe911/llama-qlora-typo" ,
        torch_dtype=torch.float16,
        device_map ="auto"
)

model = base_model

# format function
def format_output(raw_output):
    # Extract fields using regex (adjust patterns as needed)
    event = re.search(r"Event:\s*(.+)", raw_output)
    start = re.search(r"Start Time:\s*(.+)", raw_output)
    end = re.search(r"End Time:\s*(.+)", raw_output)
    date = re.search(r"Date:\s*(.+)", raw_output)
    location = re.search(r"Location:\s*(.+)", raw_output)
    duration = re.search(r"Duration:\s*(.+)", raw_output)

    # Format output nicely
    formatted = [
        f"Event: {event.group(1) if event else 'None'}",
        f"Date: {date.group(1) if date else 'None'}",
        f"Start Time: {start.group(1) if start else 'None'}",
        f"End Time: {end.group(1) if end else 'None'}",
        f"Duration: {duration.group(1) if duration else 'Unknown'}",
        f"Location: {location.group(1) if location else 'None'}",
    ]
    return "\n".join(formatted)
# Define inference function
def correct_typo(input_text):
    # Construct prompt
    prompt = f"Fix this sentence:\n{input_text}"

    # Tokenize and move to the correct device
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    # Generate output using the model
    with torch.no_grad():
        output = model.generate(**inputs, max_new_tokens=80)
    decoded = tokenizer.decode(output[0], skip_special_tokens=True)

    # Decode and return clean string
    return format_output(decoded)

# === Launch Gradio Web UI ===
gr.Interface(
    fn=correct_typo,                     # Function to call
    inputs="text",                       # Input type: plain text
    outputs="text",                      # Output type: plain text
    title=" LLaMA2 Typo Corrector by Zoe911",  # Interface title
    description="Enter messy input like '3.21 go libary at 3pm', and it returns structured event info.",
    examples=[
        "3.20 go to mall at 2pm",
        "March 30th class from 9 to 10 in Room A105"
    ]
).launch()