Update app.py

app.py

@@ -4,10 +4,7 @@ import torch
 import os
 
 # Retrieve the token from environment variables
-api_token = os.getenv("HF_TOKEN")
-if not api_token:
-    raise ValueError("HF_TOKEN environment variable is not set. Please set it to your Hugging Face API token.")
-api_token = api_token.strip()
+api_token = os.getenv("HF_TOKEN").strip()
 
 # Model name
 model_name = "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1"
@@ -36,6 +33,10 @@ def generate_response(input_text):
         # Ensure input tensor is sent to the same device as the model
         input_ids = inputs["input_ids"].to(model.device)
 
+        # Add batch dimension (if missing)
+        if len(input_ids.shape) == 1:  # If shape is (seq_len,)
+            input_ids = input_ids.unsqueeze(0)  # Add batch dimension: (1, seq_len)
+
         # Generate a response using the model
         outputs = model.generate(
             input_ids,
@@ -65,4 +66,4 @@ iface = gr.Interface(
 
 # Launch the Gradio app
 if __name__ == "__main__":
-    iface.launch(debug=True)
+    iface.launch()