Update app.py

app.py

@@ -1,7 +1,4 @@
-
-# Retrieve the token from environment variables
-#api_token = os.getenv("HF_TOKEN").strip()
-
+import os
 import torch
 from flask import Flask, request, jsonify
 from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
@@ -9,6 +6,9 @@ from PIL import Image
 import io
 import base64
 
+# Get API token from environment variable
+api_token = os.getenv("HF_TOKEN").strip()
+
 app = Flask(__name__)
 
 # Quantization configuration
@@ -19,19 +19,20 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.float16
 )
 
-# Load model
+# Load model without Flash Attention
 model = AutoModel.from_pretrained(
     "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1", 
     quantization_config=bnb_config, 
     device_map="auto", 
     torch_dtype=torch.float16, 
-    trust_remote_code=True, 
-    attn_implementation="flash_attention_2"
+    trust_remote_code=True,
+    token=api_token
 )
 
 tokenizer = AutoTokenizer.from_pretrained(
     "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1", 
-    trust_remote_code=True
+    trust_remote_code=True,
+    token=api_token
 )
 
 def decode_base64_image(base64_string):
@@ -74,7 +75,5 @@ def analyze_input():
             'message': str(e)
         }), 500
 
-
 if __name__ == '__main__':
-    app.run(debug=True)
-
+    app.run(debug=True)