use different app for testing

Dockerfile

@@ -2,10 +2,10 @@
 FROM python:3.10.9
 
 # Copy the current directory contents into the container at .
-COPY ./app ./app
+COPY . .
 
 # Set the working directory to /
-WORKDIR /trans
+WORKDIR /
 
 EXPOSE 7860
 
@@ -13,4 +13,4 @@ EXPOSE 7860
 RUN pip install --no-cache-dir --upgrade -r /requirements.txt
 
 # Start the FastAPI app on port 7860, the default port expected by Spaces
-CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,17 +1,12 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from typing import List
-import torch
+from pydantic import BaseModel
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from IndicTransToolkit import IndicProcessor
 from fastapi.middleware.cors import CORSMiddleware
 
-import os
-
-os.environ["HF_HOME"] = "/.cache"
-# Initialize FastAPI
 app = FastAPI()
 
-# Add CORS middleware
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -20,13 +15,13 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Initialize models and processors
 model = AutoModelForSeq2SeqLM.from_pretrained(
     "ai4bharat/indictrans2-en-indic-1B", trust_remote_code=True
 )
 tokenizer = AutoTokenizer.from_pretrained(
     "ai4bharat/indictrans2-en-indic-1B", trust_remote_code=True
 )
+
 ip = IndicProcessor(inference=True)
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 model = model.to(DEVICE)
@@ -58,29 +53,27 @@ def translate_text(sentences: List[str], target_lang: str):
             generated_tokens = tokenizer.batch_decode(
                 generated_tokens.detach().cpu().tolist(),
                 skip_special_tokens=True,
-                clean_up_tokenization_spaces=True,
             )
 
-        translations = ip.postprocess_batch(generated_tokens, lang=target_lang)
-        return {
-            "translations": translations,
-            "source_language": src_lang,
-            "target_language": target_lang,
-        }
+        return generated_tokens
     except Exception as e:
-        raise Exception(f"Translation failed: {str(e)}")
+        return str(e)
+
+
+@app.get("/")
+def read_root():
+    return {"Hello": "World"}
 
 
-# FastAPI routes
-@app.get("/health")
-async def health_check():
-    return {"status": "healthy"}
+class TranslateRequest(BaseModel):
+    sentences: List[str]
+    target_lang: str
 
 
-@app.post("/translate")
-async def translate_endpoint(sentences: List[str], target_lang: str):
+@app.post("/translate/")
+def translate(request: TranslateRequest):
     try:
-        result = translate_text(sentences=sentences, target_lang=target_lang)
+        result = translate_text(request.sentences, request.target_lang)
         return result
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))