Update app.py

app.py

@@ -1,52 +1,37 @@
 # app.py
 import streamlit as st
-from fastapi import FastAPI, HTTPException, Request
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel
-from typing import List
 import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from IndicTransToolkit import IndicProcessor
-import json
-from fastapi.middleware.cors import CORSMiddleware
-import uvicorn
+from typing import List
+import sys
 from starlette.applications import Starlette
-from starlette.routing import Mount, Route
+from starlette.routing import Mount
 from starlette.staticfiles import StaticFiles
-import asyncio
 import nest_asyncio
+from api import app
 
 # Enable nested event loops
 nest_asyncio.apply()
 
-# Initialize FastAPI
-app = FastAPI()
-
-# Add CORS middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    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)
+# Initialize models and processors (lazy loading)
+@st.cache_resource
+def load_models():
+    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)
+    return model, tokenizer, ip, DEVICE
 
-class TranslationRequest(BaseModel):
-    sentences: List[str]
-    target_lang: str
+# Global variables for models
+model, tokenizer, ip, DEVICE = load_models()
 
 def translate_text(sentences: List[str], target_lang: str):
     try:
@@ -90,23 +75,6 @@ def translate_text(sentences: List[str], target_lang: str):
     except Exception as e:
         raise Exception(f"Translation failed: {str(e)}")
 
-# FastAPI routes
-@app.get("/api/health")
-async def health_check():
-    return {"status": "healthy"}
-
-@app.post("/api/translate")
-async def translate_endpoint(request: TranslationRequest):
-    try:
-        result = translate_text(
-            sentences=request.sentences,
-            target_lang=request.target_lang
-        )
-        return JSONResponse(content=result)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-# Streamlit interface
 def streamlit_app():
     st.title("Indic Language Translator")
     
@@ -149,7 +117,7 @@ def streamlit_app():
     st.markdown("""
     To use the translation API, send POST requests to:
     ```
-    https://darshankr-trans-en-indic.hf.space/api/translate
+    https://YOUR-SPACE-NAME.hf.space/api/translate
     ```
     Request body format:
     ```json
@@ -163,7 +131,6 @@ def streamlit_app():
     for lang, code in target_languages.items():
         st.markdown(f"- {lang}: `{code}`")
 
-# Create a unified application
 def create_app():
     routes = [
         Mount("/api", app),
@@ -175,4 +142,5 @@ if __name__ == "__main__":
     if "streamlit" in sys.argv[0]:
         streamlit_app()
     else:
+        import uvicorn
         uvicorn.run(create_app(), host="0.0.0.0", port=7860)