Commit initial quote caster ui

Dockerfile

@@ -0,0 +1,24 @@
+FROM python:3.10-slim
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y git && rm -rf /var/lib/apt/lists/*
+
+# Set working directory
+WORKDIR /app
+
+# Copy project files
+COPY . .
+
+# Install Python packages
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Download Hugging Face model weights on build (optional to avoid slow startup)
+RUN python -c "from transformers import AutoTokenizer, AutoModel; \
+    AutoTokenizer.from_pretrained('aNameNobodyChose/quote-caster-encoder'); \
+    AutoModel.from_pretrained('aNameNobodyChose/quote-caster-encoder')"
+
+# Gradio default port
+EXPOSE 7860
+
+# Run the Gradio app
+CMD ["python", "app.py"]

README.md

@@ -1,14 +1,5 @@
----
-title: Quote Caster
-emoji: 👁
-colorFrom: pink
-colorTo: purple
-sdk: gradio
-sdk_version: 5.25.2
-app_file: app.py
-pinned: false
-license: mit
-short_description: Assigns dilogues to speakers in a story
----
+# Build quote caster gradio image
+docker build -t quote-caster-gradio .
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Run Quote caster
+docker run --rm -it -p 7860:7860 quote-caster-gradio

app.py

@@ -0,0 +1,66 @@
+import gradio as gr
+from transformers import AutoModel, AutoTokenizer
+from sklearn.cluster import KMeans
+from kneed import KneeLocator
+import torch
+import json
+
+def encode_quote(context: str, dialogue: str, tokenizer, model) -> torch.Tensor:
+    """
+    Encode a single quote using [CLS] token from BERT.
+    """
+    text = f"{context} [SEP] {dialogue}"
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        padding=True,
+        max_length=512
+    )
+    outputs = model(**inputs)
+    cls_embedding = outputs.last_hidden_state[:, 0, :]  # [CLS] token
+    return cls_embedding.squeeze(0)
+
+def load_encoder():
+    tokenizer = AutoTokenizer.from_pretrained("aNameNobodyChose/quote-caster-encoder")
+    model = AutoModel.from_pretrained("aNameNobodyChose/quote-caster-encoder")
+    model.eval()
+    return tokenizer, model
+
+def embed_quotes(data, tokenizer, model):
+    embeddings = []
+    for ex in data:
+        emb = encode_quote(ex["context"], ex["quote"], tokenizer, model)
+        embeddings.append(emb)
+    return torch.stack(embeddings)
+
+def auto_k_via_elbow(embeddings, max_k=10):
+    X = embeddings.detach().numpy()
+    inertias = []
+    for k in range(1, max_k + 1):
+        kmeans = KMeans(n_clusters=k, random_state=42, n_init='auto')
+        kmeans.fit(X)
+        inertias.append(kmeans.inertia_)
+    knee = KneeLocator(range(1, max_k + 1), inertias, curve="convex", direction="decreasing")
+    return knee.knee or 2
+
+def predict(input_text):
+    try:
+        data = json.loads(input_text)
+        tokenizer, model = load_encoder()
+        embeddings = embed_quotes(data, tokenizer, model)
+        k = auto_k_via_elbow(embeddings)
+        labels = KMeans(n_clusters=k).fit_predict(embeddings.detach().numpy())
+        for quote, cluster_id in zip(data, labels):
+            quote["predicted_speaker"] = f"SPEAKER_{cluster_id}"
+        return json.dumps(data, indent=2, ensure_ascii=False)
+    except Exception as e:
+        return f"❌ Error: {e}"
+
+gr.Interface(
+    fn=predict,
+    inputs=gr.Textbox(lines=20, label="Paste quote-context JSON"),
+    outputs="textbox",
+    title="🗣️ QuoteCaster - Speaker Attribution from Quotes",
+    description="Paste a list of quotes with their context and get clustered speaker predictions using a transformer-based model."
+).launch(server_name="0.0.0.0", server_port=7860)

requirements.txt

@@ -0,0 +1,5 @@
+gradio
+torch
+transformers==4.39.3
+scikit-learn==1.3.2
+kneed

space.yaml

@@ -0,0 +1 @@
+sdk: docker