Update app.py

app.py

@@ -1,58 +1,61 @@
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-from dotenv import load_dotenv
-from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-import os
-
-
-load_dotenv()
-
-
-os.environ["HF_TOKEN"] = os.getenv('HF_TOKEN')
-
-
-app = FastAPI()
-
-
-origins = ["*"]
-
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=origins,
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-
-class LanguageTextModel(BaseModel):
-    languageText: str
-    sourceLanguageCode: str
-    targetLanguageCode: str
-
-
-@app.get('/')
-def welcome():
-    return {
-        'success': True,
-        'message': 'server of "nllb language translator" is up and running successfully '
-    }
-
-
-@app.post('/translate')
-async def translate_text(input: LanguageTextModel):
-    
-    model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
-
-    tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
-
-    translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=input.sourceLanguageCode, tgt_lang=input.targetLanguageCode, max_length=400)
-
-    response = translator('I am not feeling well')
-
-    return {
-        "success": True,
-        "translated_text": response[0]['translation_text']
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from dotenv import load_dotenv
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import os
+
+
+load_dotenv()
+
+
+os.environ["HF_TOKEN"] = os.getenv('HF_TOKEN')
+
+os.environ["HF_HOME"] = "/code/.cache/huggingface"
+
+
+
+app = FastAPI()
+
+
+origins = ["*"]
+
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+class LanguageTextModel(BaseModel):
+    languageText: str
+    sourceLanguageCode: str
+    targetLanguageCode: str
+
+
+@app.get('/')
+def welcome():
+    return {
+        'success': True,
+        'message': 'server of "nllb language translator" is up and running successfully '
+    }
+
+
+@app.post('/translate')
+async def translate_text(input: LanguageTextModel):
+    
+    model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
+
+    tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
+
+    translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=input.sourceLanguageCode, tgt_lang=input.targetLanguageCode, max_length=400)
+
+    response = translator('I am not feeling well')
+
+    return {
+        "success": True,
+        "translated_text": response[0]['translation_text']
     }