Update app.py

app.py

@@ -14,24 +14,24 @@ import torch
 import cv2
 from gradio_client import Client, file
 
+# Function to generate an image using another model
 def image_gen(prompt):
     client = Client("KingNish/Image-Gen-Pro")
-    return client.predict("Image Generation",None, prompt, api_name="/image_gen_pro")
+    return client.predict("Image Generation", None, prompt, api_name="/image_gen_pro")
 
+# Load the processor and model for image-based QnA (LLaVA model)
 model_id = "llava-hf/llava-interleave-qwen-0.5b-hf"
-
 processor = LlavaProcessor.from_pretrained(model_id)
-
 model = LlavaForConditionalGeneration.from_pretrained(model_id)
 model.to("cpu")
 
-
+# Function to process images with text input
 def llava(message, history):
     if message["files"]:
         image = message["files"][0]    
     else:
         for hist in history:
-            if type(hist[0])==tuple:
+            if type(hist[0]) == tuple:
                 image = hist[0][0]
         
     txt = message["text"]
@@ -43,12 +43,14 @@ def llava(message, history):
     inputs = processor(prompt, image, return_tensors="pt")
     return inputs
 
+# Helper function to extract text from a webpage
 def extract_text_from_webpage(html_content):
     soup = BeautifulSoup(html_content, 'html.parser')
     for tag in soup(["script", "style", "header", "footer"]):
         tag.extract()
     return soup.get_text(strip=True)
 
+# Function to search the web using Google
 def search(query):
     term = query
     start = 0
@@ -88,8 +90,8 @@ client_yi = InferenceClient("01-ai/Yi-1.5-34B-Chat")
 # Define the main chat function
 def respond(message, history):
     func_caller = []
-
     user_prompt = message
+
     # Handle image processing
     if message["files"]:
         inputs = llava(message, history)
@@ -101,9 +103,11 @@ def respond(message, history):
     
         buffer = ""
         for new_text in streamer:
-            buffer += new_text
-            yield buffer
+            if new_text not in ["<|im_end|>", "<|endoftext|>"]:  # Ignore special tokens
+                buffer += new_text
+                yield buffer
     else:
+        # Functions metadata for invoking different models or functions
         functions_metadata = [
             {"type": "function", "function": {"name": "web_search", "description": "Search query on google", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "web search query"}}, "required": ["query"]}}},
             {"type": "function", "function": {"name": "general_query", "description": "Reply general query of USER", "parameters": {"type": "object", "properties": {"prompt": {"type": "string", "description": "A detailed prompt"}}, "required": ["prompt"]}}},
@@ -120,45 +124,41 @@ def respond(message, history):
     
         response = client_gemma.chat_completion(func_caller, max_tokens=200)
         response = str(response)
+        
+        # Filtering and processing response
         try:
             response = response[int(response.find("{")):int(response.rindex("</"))]
         except:
             response = response[int(response.find("{")):(int(response.rfind("}"))+1)]
-        response = response.replace("\\n", "")
-        response = response.replace("\\'", "'")
-        response = response.replace('\\"', '"')
-        response = response.replace('\\', '')
+        response = response.replace("\\n", "").replace("\\'", "'").replace('\\"', '"').replace('\\', '')
         print(f"\n{response}")
     
         try:
             json_data = json.loads(str(response))
             if json_data["name"] == "web_search":
                 query = json_data["arguments"]["query"]
-               # gr.Info("Searching Web")
                 web_results = search(query)
-               # gr.Info("Extracting relevant Info")
                 web2 = ' '.join([f"Link: {res['link']}\nText: {res['text']}\n\n" for res in web_results])
                 messages = f"<|im_start|>system\n Hi 👋, I am Nora,mini a helpful assistant.Ask me! I will do my best!! <|im_end|>"
                 for msg in history:
                     messages += f"\n<|im_start|>user\n{str(msg[0])}<|im_end|>"
                     messages += f"\n<|im_start|>assistant\n{str(msg[1])}<|im_end|>"
-                messages+=f"\n<|im_start|>user\n{message_text}<|im_end|>\n<|im_start|>web_result\n{web2}<|im_end|>\n<|im_start|>assistant\n"
+                messages += f"\n<|im_start|>user\n{message_text}<|im_end|>\n<|im_start|>web_result\n{web2}<|im_end|>\n<|im_start|>assistant\n"
                 stream = client_mixtral.text_generation(messages, max_new_tokens=2000, do_sample=True, stream=True, details=True, return_full_text=False)
                 output = ""
                 for response in stream:
-                    if not response.token.text == "<|im_end|>":
+                    if not response.token.text in ["<|im_end|>", "<|endoftext|>"]:  # Exclude special tokens
                         output += response.token.text
                         yield output
             elif json_data["name"] == "image_generation":
                 query = json_data["arguments"]["query"]
-                gr.Info("Generating Image, Please wait 10 sec...")
                 yield "Generating Image, Please wait 10 sec..."
                 try:
                     image = image_gen(f"{str(query)}")
                     yield gr.Image(image[1])
                 except:
                     client_sd3 = InferenceClient("stabilityai/stable-diffusion-3-medium-diffusers")
-                    seed = random.randint(0,999999)
+                    seed = random.randint(0, 999999)
                     image = client_sd3.text_to_image(query, negative_prompt=f"{seed}")
                     yield gr.Image(image)
             elif json_data["name"] == "image_qna":
@@ -168,33 +168,35 @@ def respond(message, history):
 
                 thread = Thread(target=model.generate, kwargs=generation_kwargs)
                 thread.start()
-    
+
                 buffer = ""
                 for new_text in streamer:
-                    buffer += new_text
-                    yield buffer
+                    if new_text not in ["<|im_end|>", "<|endoftext|>"]:  # Ignore special tokens
+                        buffer += new_text
+                        yield buffer
             else:
                 messages = f"<|im_start|>system\n 👋, I am Nora,mini a helpful assistant.Ask me! I will do my best!!<|im_end|>"
                 for msg in history:
                     messages += f"\n<|im_start|>user\n{str(msg[0])}<|im_end|>"
                     messages += f"\n<|im_start|>assistant\n{str(msg[1])}<|im_end|>"
-                messages+=f"\n<|im_start|>user\n{message_text}<|im_end|>\n<|im_start|>assistant\n"
+                messages += f"\n<|im_start|>user\n{message_text}<|im_end|>\n<|im_start|>assistant\n"
                 stream = client_yi.text_generation(messages, max_new_tokens=2000, do_sample=True, stream=True, details=True, return_full_text=False)
                 output = ""
                 for response in stream:
-                    if not response.token.text == "<|endoftext|>":
+                    if response.token.text not in ["<|im_end|>", "<|endoftext|>"]:  # Ignore special tokens
                         output += response.token.text
                         yield output
         except:
-            messages = f"<|start_header_id|>system\nHi 👋, I am Nora,mini a helpful assistant.Ask me! I will do my best!!<|end_header_id|>"
+            # Handle the case where JSON parsing or function calling fails
+            messages = f"<|im_start|>system\nHi 👋, I am Nora,mini a helpful assistant.Ask me! I will do my best!!<|im_end|>"
             for msg in history:
-                messages += f"\n<|start_header_id|>user\n{str(msg[0])}<|end_header_id|>"
-                messages += f"\n<|start_header_id|>assistant\n{str(msg[1])}<|end_header_id|>"
-            messages+=f"\n<|start_header_id|>user\n{message_text}<|end_header_id|>\n<|start_header_id|>assistant\n"
+                messages += f"\n<|im_start|>user\n{str(msg[0])}<|im_end|>"
+                messages += f"\n<|im_start|>assistant\n{str(msg[1])}<|im_end|>"
+            messages += f"\n<|im_start|>user\n{message_text}<|im_end|>\n<|im_start|>assistant\n"
             stream = client_llama.text_generation(messages, max_new_tokens=2000, do_sample=True, stream=True, details=True, return_full_text=False)
             output = ""
             for response in stream:
-                if not response.token.text == "<|eot_id|>":
+                if response.token.text not in ["<|eot_id|>", "<|im_end|>"]:  # Ignore special tokens
                     output += response.token.text
                     yield output
 
@@ -205,6 +207,9 @@ demo = gr.ChatInterface(
     textbox=gr.MultimodalTextbox(),
     multimodal=True,
     concurrency_limit=200,
-    cache_examples=False,css="footer{display:none !important}"
+    cache_examples=False,
+    css="footer{display:none !important}"
 )
+
+# Launch the Gradio app
 demo.launch()