Update app.py

app.py

@@ -45,7 +45,7 @@ text_tokenizer = AutoTokenizer.from_pretrained(
 with torch.no_grad():
     tuned_phi2 = AutoModelForCausalLM.from_pretrained(
         "stage2_adaptor", trust_remote_code=True,
-        device=device, torch_dtype=torch.float16
+        device_map="auto",
     )
     base_phi2_text = AutoModelForCausalLM.from_pretrained(
         phi_model_name, trust_remote_code=True,
@@ -93,12 +93,14 @@ def process_audio(audio):
 
 def generate_response(image, audio, text, count):
     count = int(count)
+    
     if audio:
         text_from_audio = process_audio(audio)
     if text:
         overall_input = text + text_from_audio
     if image:
         img_tokens = process_image(image)
+        overall_input = "Question: " + overall_input + "Answer:" 
         q_tokens = text_tokenizer.encode(
             overall_input,
             return_tensors='pt').to(device)
@@ -120,10 +122,7 @@ def generate_response(image, audio, text, count):
     else:
         return process_text(overall_input, count)
 
-    return prediction[0].strip('<|endoftext|>').rstrip("\n")
 
-
-%%blocks
 with gr.Blocks() as demo:
     gr.Markdown("# **AnyModeAssistant**")
     gr.Markdown("Use any mode text/image/audio to interact with AI assistant")