Spaces:
Running on Zero
Running on Zero
File size: 2,834 Bytes
88a619b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """File I/O: reading source documents (txt/docx/pdf) and writing translated
output (txt/docx/json). Segments are plain dicts: {"source": str, "target": str}.
"""
import json
from pathlib import Path
from .chunking import split_into_segments
def _extract_txt(path: str) -> str:
return Path(path).read_text(encoding="utf-8", errors="replace")
def _extract_docx(path: str) -> str:
import docx
doc = docx.Document(path)
return "\n\n".join(p.text for p in doc.paragraphs if p.text.strip())
def _extract_pdf(path: str) -> str:
from pypdf import PdfReader
reader = PdfReader(path)
return "\n\n".join(page.extract_text() or "" for page in reader.pages)
_EXTRACTORS = {
".txt": _extract_txt,
".docx": _extract_docx,
".pdf": _extract_pdf,
}
def load_source_file(path: str, max_chars: int = 400) -> list[dict]:
"""Extract text from a .txt/.docx/.pdf file and split it into segments."""
ext = Path(path).suffix.lower()
extractor = _EXTRACTORS.get(ext)
if extractor is None:
raise ValueError(f"Unsupported file type: {ext}")
text = extractor(path)
if not text.strip():
raise ValueError("No text could be extracted from this file.")
chunks = split_into_segments(text, max_chars=max_chars)
return [{"source": c, "target": ""} for c in chunks]
def parse_json(file_content: str | bytes) -> list[dict]:
"""Parse a previously-exported JSON file back into segments (for resume)."""
if isinstance(file_content, bytes):
file_content = file_content.decode("utf-8")
data = json.loads(file_content)
if not isinstance(data, list) or not data:
raise ValueError("JSON must be a non-empty list of objects.")
segments = []
for i, item in enumerate(data):
if not isinstance(item, dict) or "source" not in item:
raise ValueError(f"Item {i} is missing 'source'.")
segments.append({"source": item.get("source", ""), "target": item.get("target", "")})
return segments
def export_to_txt(segments: list[dict], which: str = "target") -> str:
"""Join segment text with blank lines. which='target' or 'source'."""
texts = [seg.get(which, "").strip() for seg in segments]
texts = [t for t in texts if t]
return "\n\n".join(texts) + "\n" if texts else ""
def export_to_docx(segments: list[dict], which: str = "target") -> bytes:
import io
import docx
doc = docx.Document()
for seg in segments:
text = seg.get(which, "").strip()
if text:
doc.add_paragraph(text)
buf = io.BytesIO()
doc.save(buf)
return buf.getvalue()
def export_to_json(segments: list[dict]) -> str:
return json.dumps(
[{"source": s.get("source", ""), "target": s.get("target", "")} for s in segments],
ensure_ascii=False, indent=2,
)
|