Update app.py

app.py

@@ -1,44 +1,52 @@
-from fastapi import FastAPI, UploadFile, File
-from transformers import MarianMTModel, MarianTokenizer
+from fastapi import FastAPI, File, UploadFile
+from transformers import pipeline
 import PyPDF2
 import docx
-import io
+import os
+import uvicorn
+from io import BytesIO
 
 app = FastAPI()
 
-# Charger le modèle MarianMT pour la traduction (ex: anglais → français)
-MODEL_NAME = "Helsinki-NLP/opus-mt-en-fr"
-tokenizer = MarianTokenizer.from_pretrained(MODEL_NAME)
-model = MarianMTModel.from_pretrained(MODEL_NAME)
+# Charger le modèle de traduction depuis Hugging Face (Exemple : Anglais -> Français)
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-en-fr")
 
-def translate_text(text, src_lang="en", tgt_lang="fr"):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-    translated = model.generate(**inputs)
-    return tokenizer.batch_decode(translated, skip_special_tokens=True)[0]
-
-def extract_text_from_pdf(pdf_file):
-    reader = PyPDF2.PdfReader(pdf_file)
-    text = " ".join(page.extract_text() for page in reader.pages if page.extract_text())
+def extract_text_from_pdf(file: BytesIO) -> str:
+    """Extrait le texte d'un fichier PDF."""
+    reader = PyPDF2.PdfReader(file)
+    text = "\n".join([page.extract_text() for page in reader.pages if page.extract_text()])
     return text
 
-def extract_text_from_docx(docx_file):
-    doc = docx.Document(docx_file)
-    text = " ".join([p.text for p in doc.paragraphs])
+def extract_text_from_docx(file: BytesIO) -> str:
+    """Extrait le texte d'un fichier DOCX."""
+    doc = docx.Document(file)
+    text = "\n".join([para.text for para in doc.paragraphs])
     return text
 
 @app.post("/translate/")
-async def translate_file(file: UploadFile = File(...), src_lang: str = "en", tgt_lang: str = "fr"):
-    if file.filename.endswith(".pdf"):
-        text = extract_text_from_pdf(io.BytesIO(await file.read()))
-    elif file.filename.endswith(".docx"):
-        text = extract_text_from_docx(io.BytesIO(await file.read()))
-    else:
-        return {"error": "Format non supporté. Utilise PDF ou DOCX."}
-    
-    translated_text = translate_text(text, src_lang, tgt_lang)
-    return {"translated_text": translated_text}
-
-@app.get("/")
-def home():
-    return {"message": "Bienvenue sur l'API de traduction de fichiers !"}
+async def translate_file(file: UploadFile = File(...)):
+    """Endpoint pour traduire un fichier PDF ou DOCX."""
+    try:
+        contents = await file.read()
+        file_io = BytesIO(contents)
+        file_extension = file.filename.split(".")[-1]
+
+        if file_extension == "pdf":
+            text = extract_text_from_pdf(file_io)
+        elif file_extension == "docx":
+            text = extract_text_from_docx(file_io)
+        else:
+            return {"error": "Format non supporté. Utilisez PDF ou DOCX."}
+
+        # Traduire le texte
+        translation = translator(text, max_length=1000)
+        translated_text = " ".join([t["translation_text"] for t in translation])
+
+        return {"original_text": text[:500], "translated_text": translated_text[:500]}  # Limite pour affichage
+    except Exception as e:
+        return {"error": str(e)}
+
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 7860))
+    uvicorn.run(app, host="0.0.0.0", port=port)