Update app.py

app.py

@@ -6,27 +6,27 @@ HUGGING_FACE_USER_NAME = "elalimy"
 model_name = "my_awesome_peft_finetuned_helsinki_model"
 peft_model_id = f"{HUGGING_FACE_USER_NAME}/{model_name}"
 
-# Flask app class
-app = Flask(__name__, template_folder='templates')  # Specify the templates folder
+# Load model configuration (assuming it's saved locally)
+config = PeftConfig.from_pretrained(peft_model_id)
+# Load the base model from its local directory (replace with actual model type)
+base_model = AutoModelForSeq2SeqLM.from_pretrained(config.base_model_name_or_path, return_dict=True, load_in_8bit=False)
 
+# Load the tokenizer from its local directory (replace with actual tokenizer type)
+tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 
-def generate_translation(source_text, device="cpu"):
-    # Load model configuration (assuming it's saved locally)
-    config = PeftConfig.from_pretrained(peft_model_id)
-    # Load the base model from its local directory (replace with actual model type)
-    model = AutoModelForSeq2SeqLM.from_pretrained(config.base_model_name_or_path, return_dict=True, load_in_8bit=False)
+# Load the Peft model (assuming it's a custom class or adaptation)
+AI_model = PeftModel.from_pretrained(base_model, peft_model_id)
 
-    # Load the tokenizer from its local directory (replace with actual tokenizer type)
-    tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
+# Flask app class
+app = Flask(__name__, template_folder='templates')  # Specify the templates folder
 
-    # Load the Peft model (assuming it's a custom class or adaptation)
-    AI_model = PeftModel.from_pretrained(model, peft_model_id)
 
+def generate_translation(source_text, device="cpu"):
     # Encode the source text
     input_ids = tokenizer.encode(source_text, return_tensors='pt').to(device)
 
     # Move the model to the same device as input_ids
-    model = model.to(device)
+    model = base_model.to(device)
 
     # Generate the translation with adjusted decoding parameters
     generated_ids = model.generate(
@@ -62,4 +62,4 @@ def translate_text():
 
 
 if __name__ == "__main__":
-    app.run()
+    app.run()