Spaces:
Sleeping
Sleeping
File size: 6,289 Bytes
d333245 4b69e7c d333245 4b69e7c d333245 774c080 d333245 774c080 d333245 4b69e7c d333245 4b69e7c d333245 |
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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
import gradio as gr
from transformers import pipeline, MarianTokenizer, AutoModelForSeq2SeqLM
import torch
import unicodedata
import re
import whisper
import tempfile
import os
import nltk
nltk.download('punkt')
from nltk.tokenize import sent_tokenize
# Additions for file processing
import fitz # PyMuPDF for PDF
import docx
from bs4 import BeautifulSoup
import markdown2
import chardet
# --- Device selection ---
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# --- Load translation models ---
def load_models():
en_dar_model_path = "/LocaleNLP/English_Hausa"
en_wol_model_path = "/LocaleNLP/eng_wolof"
en_hau_model_path = "/LocaleNLP/English_Darija"
en_dar_model = AutoModelForSeq2SeqLM.from_pretrained(en_dar_model_path).to(device)
en_dar_tokenizer = MarianTokenizer.from_pretrained(en_dar_model_path)
en_wol_model = AutoModelForSeq2SeqLM.from_pretrained(en_wol_model_path).to(device)
en_wol_tokenizer = MarianTokenizer.from_pretrained(en_wol_model_path)
en_hau_model = AutoModelForSeq2SeqLM.from_pretrained(en_hau_model_path).to(device)
en_hau_tokenizer = MarianTokenizer.from_pretrained(en_hau_model_path)
en_dar_translator = pipeline("translation", model=en_dar_model, tokenizer=en_dar_tokenizer, device=0 if device.type == 'cuda' else -1)
en_wol_translator = pipeline("translation", model=en_wol_model, tokenizer=en_wol_tokenizer, device=0 if device.type == 'cuda' else -1)
en_hau_translator = pipeline("translation", model=en_hau_model, tokenizer=en_hau_tokenizer, device=0 if device.type == 'cuda' else -1)
return en_dar_translator, en_hau_translator, en_wol_translator
def load_whisper_model():
return whisper.load_model("base")
def transcribe_audio(audio_path):
whisper_model = load_whisper_model()
return whisper_model.transcribe(audio_path)["text"]
def translate(text, target_lang):
en_dar_translator, en_hau_translator, en_wol_translator = load_models()
if target_lang == "Darija (Morocco)":
translator = en_dar_translator
elif target_lang == "Hausa (Nigeria)":
translator = en_hau_translator
elif target_lang == "Wolof (Senegal)":
translator = en_wol_translator
else:
raise ValueError("Unsupported target language")
lang_tag = {
"Darija (Morocco)": ">>dar<<",
"Hausa (Nigeria)": ">>hau<<",
"Wolof (Senegal)": ">>wol<<"
}[target_lang]
paragraphs = text.split("\n")
translated_output = []
with torch.no_grad():
for para in paragraphs:
if not para.strip():
translated_output.append("")
continue
sentences = [s.strip() for s in para.split('. ') if s.strip()]
formatted = [f"{lang_tag} {s}" for s in sentences]
results = translator(formatted,
max_length=5000,
num_beams=5,
early_stopping=True,
no_repeat_ngram_size=3,
repetition_penalty=1.5,
length_penalty=1.2)
translated_sentences = [r['translation_text'].capitalize() for r in results]
translated_output.append('. '.join(translated_sentences))
return "\n".join(translated_output)
# --- Extract text from file ---
def extract_text_from_file(uploaded_file):
file_type = uploaded_file.name.split('.')[-1].lower()
content = uploaded_file.read()
if file_type == "pdf":
with fitz.open(stream=content, filetype="pdf") as doc:
return "\n".join([page.get_text() for page in doc])
elif file_type == "docx":
doc = docx.Document(uploaded_file)
return "\n".join([para.text for para in doc.paragraphs])
else:
encoding = chardet.detect(content)['encoding']
if encoding:
content = content.decode(encoding, errors='ignore')
if file_type in ("html", "htm"):
soup = BeautifulSoup(content, "html.parser")
return soup.get_text()
elif file_type == "md":
html = markdown2.markdown(content)
soup = BeautifulSoup(html, "html.parser")
return soup.get_text()
elif file_type == "srt":
return re.sub(r"\d+\n\d{2}:\d{2}:\d{2},\d{3} --> .*?\n", "", content)
elif file_type in ("txt", "text"):
return content
else:
raise ValueError("Unsupported file type")
# --- Main Gradio Function ---
def process(input_mode, target_lang, text_input, audio_input, file_input):
input_text = ""
if input_mode == "Text" and text_input:
input_text = text_input
elif input_mode == "Audio" and audio_input:
with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
tmp.write(audio_input.read())
tmp_path = tmp.name
input_text = transcribe_audio(audio_input)
os.remove(tmp_path)
elif input_mode == "File" and file_input:
input_text = extract_text_from_file(file_input)
if not input_text.strip():
return "", "No input text provided."
translated_text = translate(input_text, target_lang)
return input_text, translated_text
# --- Gradio Interface ---
with gr.Blocks() as demo:
gr.Markdown("## 🌐 LocaleNLP Translator — English ↔ Darija / Hausa / Wolof")
with gr.Row():
input_mode = gr.Dropdown(["Text", "Audio", "File"], label="Select input mode")
target_lang = gr.Dropdown(["Darija (Morocco)", "Hausa (Nigeria)", "Wolof (Senegal)"], label="Select target language")
with gr.Row():
text_input = gr.Textbox(label="Enter English text", lines=10)
audio_input = gr.Audio(type="filepath", label="Upload Audio")
file_input = gr.File(label="Upload Document")
with gr.Row():
extracted_text = gr.Textbox(label="Extracted / Transcribed Text", lines=10)
translated_output = gr.Textbox(label="Translated Text", lines=10)
run_btn = gr.Button("Translate")
run_btn.click(process, inputs=[input_mode, target_lang, text_input, audio_input, file_input], outputs=[extracted_text, translated_output])
if __name__ == "__main__":
demo.launch()
|