import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch

# ----------------------------
# Model setup (lightweight & free)
# ----------------------------
MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.float32,
    device_map="auto"
)

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer
)

# ----------------------------
# IELTS Scoring Function
# ----------------------------
def score_ielts_writing(text):
    prompt = f"""
You are a certified IELTS Writing examiner.

Ignore the essay content in your output.
Score the essay strictly using IELTS band descriptors.

Output ONLY the final scores in JSON format.
Do NOT repeat the essay.

Score these criteria:
- Task Response (TR)
- Coherence and Cohesion (CC)
- Lexical Resource (LR)
- Grammatical Range and Accuracy (GRA)

MANDATORY JSON FORMAT:
{{
  "task_response": 0.0,
  "coherence_cohesion": 0.0,
  "lexical_resource": 0.0,
  "grammatical_range_accuracy": 0.0,
  "overall_band": 0.0,
  "feedback": "2–3 sentence examiner feedback"
}}

ESSAY (QUESTION + ESSAY):
{text}

Now, output ONLY the JSON scores as above.
"""

    output = pipe(
        prompt,
        max_new_tokens=450,  # enough for 300-word essay + JSON
        temperature=0.3,     # deterministic scoring
        do_sample=False,
        repetition_penalty=1.05
    )[0]["generated_text"]

    # remove the prompt part from the output
    return output[len(prompt):].strip()

# ----------------------------
# Gradio UI
# ----------------------------
iface = gr.Interface(
    fn=score_ielts_writing,
    inputs=gr.Textbox(
        lines=18,
        label="Paste IELTS QUESTION + ESSAY",
        placeholder="QUESTION:\n...\n\nESSAY:\n..."
    ),
    outputs=gr.Textbox(label="IELTS Band Scores (JSON)"),
    title="IELTS Writing AI Scorer (Free • CPU • Fast)",
    description="Paste the IELTS Writing Task 2 question followed by the essay. Returns band scores + examiner feedback."
)

iface.launch()