Update app.py

app.py

@@ -2,6 +2,9 @@ import torch
 import torch.nn as nn
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from torchvision import transforms
+from transformers import CLIPProcessor, CLIPModel
+from PIL import Image
 
 
 class _MLPVectorProjector(nn.Module):
@@ -29,17 +32,36 @@ tokenizer_text = AutoTokenizer.from_pretrained(model_name, trust_remote_code=Tru
 
 ## Audio model
 model_name_audio = "openai/whisper-small"
-#processor_audio = WhisperProcessor.from_pretrained("openai/whisper-small")
-#model_audio = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
-#model_audio.config.forced_decoder_ids = None
-pipe = pipeline(
-    task="automatic-speech-recognition",
-    model=model_name_audio,
-    chunk_length_s=30,
-    device="cpu",
-)
+pipe = pipeline(task="automatic-speech-recognition", model=model_name_audio, 
+    chunk_length_s=30, device="cpu",)
 
 ## image model
+#Clip model
+model_id_clip = "openai/clip-vit-base-patch16"
+model_clip = CLIPModel.from_pretrained(model_id_clip).to("cpu")
+processor_clip = CLIPProcessor.from_pretrained(model_id_clip)
+
+# Preprocess the image for clip
+def preprocess_image(image_path):
+    image = Image.open(image_path).convert("RGB")
+    image = transforms.Resize((224, 224))(image)
+    image = transforms.ToTensor()(image)
+    return image.unsqueeze(0)
+
+# Get clip encoding
+def encode_image(image_path):
+    image = preprocess_image(image_path).to("cpu")    
+    # Dummy input_ids for text
+    dummy_text = ""
+    inputs = processor_clip(text=dummy_text, images=image, return_tensors="pt", padding=True)    
+    outputs = model_clip(**inputs)
+    img_embedding = outputs.image_embeds
+    return img_embedding  
+
+#Get the projection model
+
+#Get the fine-tuned phi-2 model
+    
 
 def example_inference(input_text, count): #, image, img_qn, audio):
     pred_text = textMode(input_text, count)
@@ -54,9 +76,9 @@ def textMode(text, count):
     phi2_text.generate(
         **inputs, 
         max_new_tokens=count,
-        bos_token_id=tokenizer.bos_token_id, 
-        eos_token_id=tokenizer.eos_token_id,
-        pad_token_id=tokenizer.pad_token_id
+        bos_token_id=tokenizer_text.bos_token_id, 
+        eos_token_id=tokenizer_text.eos_token_id,
+        pad_token_id=tokenizer_text.pad_token_id
     )
     )
     return prediction[0].rstrip('<|endoftext|>').rstrip("\n")
@@ -64,6 +86,7 @@ def textMode(text, count):
 
 
 def imageMode(image, question):
+    image_embedding = encode_image(image)
     return "In progress"
 
 def audioMode(audio):