Update app.py

app.py

@@ -9,14 +9,16 @@ import random
 from transformers import LlavaProcessor, LlavaForConditionalGeneration, TextIteratorStreamer
 from threading import Thread
 import re
-import time 
+import time
 import torch
 import cv2
 from gradio_client import Client, file
 
+
 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")
+
 
 model_id = "llava-hf/llava-interleave-qwen-0.5b-hf"
 
@@ -28,27 +30,29 @@ model.to("cpu")
 
 def llava(message, history):
     if message["files"]:
-        image = message["files"][0]    
+        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"]
-        
+
     gr.Info("Analyzing image")
     image = Image.open(image).convert("RGB")
     prompt = f"<|im_start|>user <image>\n{txt}<|im_start|>assistant"
-    
+
     inputs = processor(prompt, image, return_tensors="pt")
     return inputs
 
+
 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)
 
+
 def search(query):
     term = query
     start = 0
@@ -69,7 +73,9 @@ def search(query):
             link = result.find("a", href=True)
             link = link["href"]
             try:
-                webpage = session.get(link, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"}, timeout=5, verify=False)
+                webpage = session.get(link, headers={
+                    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"},
+                                      timeout=5, verify=False)
                 webpage.raise_for_status()
                 visible_text = extract_text_from_webpage(webpage.text)
                 if len(visible_text) > max_chars_per_page:
@@ -79,12 +85,14 @@ def search(query):
                 all_results.append({"link": link, "text": None})
     return all_results
 
+
 # Initialize inference clients for different models
 client_gemma = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
 client_mixtral = InferenceClient("NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO")
 client_llama = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct")
 client_yi = InferenceClient("01-ai/Yi-1.5-34B-Chat")
 
+
 # Define the main chat function
 def respond(message, history):
     func_caller = []
@@ -98,17 +106,31 @@ 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
     else:
         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"]}}},
-            {"type": "function", "function": {"name": "image_generation", "description": "Generate image for user", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "image generation prompt"}}, "required": ["query"]}}},
-            {"type": "function", "function": {"name": "image_qna", "description": "Answer question asked by user related to image", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "Question by user"}}, "required": ["query"]}}},
+            {"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"]}}},
+            {"type": "function", "function": {"name": "image_generation", "description": "Generate image for user",
+                                              "parameters": {"type": "object", "properties": {
+                                                  "query": {"type": "string",
+                                                            "description": "image generation prompt"}},
+                                                             "required": ["query"]}}},
+            {"type": "function",
+             "function": {"name": "image_qna", "description": "Answer question asked by user related to image",
+                          "parameters": {"type": "object",
+                                         "properties": {"query": {"type": "string", "description": "Question by user"}},
+                                         "required": ["query"]}}},
         ]
 
         for msg in history:
@@ -116,20 +138,21 @@ def respond(message, history):
             func_caller.append({"role": "assistant", "content": f"{str(msg[1])}"})
 
         message_text = message["text"]
-        func_caller.append({"role": "user", "content": f'[SYSTEM]You are a helpful assistant. You have access to the following functions: \n {str(functions_metadata)}\n\nTo use these functions respond with:\n<functioncall> {{ "name": "function_name", "arguments": {{ "arg_1": "value_1", "arg_1": "value_1", ... }} }}  </functioncall>  [USER] {message_text}'})
-    
+        func_caller.append({"role": "user",
+                            "content": f'[SYSTEM]You are a helpful assistant. You have access to the following functions: \n {str(functions_metadata)}\n\nTo use these functions respond with:\n<functioncall> {{ "name": "function_name", "arguments": {{ "arg_1": "value_1", "arg_1": "value_1", ... }} }}  </functioncall>  [USER] {message_text}'})
+
         response = client_gemma.chat_completion(func_caller, max_tokens=200)
         response = str(response)
         try:
             response = response[int(response.find("{")):int(response.rindex("</"))]
         except:
-            response = response[int(response.find("{")):(int(response.rfind("}"))+1)]
+            response = response[int(response.find("{")):(int(response.rfind("}")) + 1)]
         response = response.replace("\\n", "")
         response = response.replace("\\'", "'")
         response = response.replace('\\"', '"')
         response = response.replace('\\', '')
         print(f"\n{response}")
-    
+
         try:
             json_data = json.loads(str(response))
             if json_data["name"] == "web_search":
@@ -142,11 +165,12 @@ def respond(message, history):
                 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}\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)
+                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 == "hello":
                         output += response.token.text
                         yield output
             elif json_data["name"] == "image_generation":
@@ -158,7 +182,7 @@ def respond(message, history):
                     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,7 +192,7 @@ def respond(message, history):
 
                 thread = Thread(target=model.generate, kwargs=generation_kwargs)
                 thread.start()
-    
+
                 buffer = ""
                 for new_text in streamer:
                     buffer += new_text
@@ -178,8 +202,9 @@ def respond(message, history):
                 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"
-                stream = client_yi.text_generation(messages, max_new_tokens=2000, do_sample=True, stream=True, details=True, return_full_text=False)
+                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|>":
@@ -190,19 +215,21 @@ def respond(message, history):
             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"
-            stream = client_llama.text_generation(messages, max_new_tokens=2000, do_sample=True, stream=True, details=True, return_full_text=False)
+            messages += f"\n<|start_header_id|>user\n{message_text}<|end_header_id|>\n<|start_header_id|>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|>":
                     output += response.token.text
                     yield output
 
+
 # Create the Gradio interface
 demo = gr.ChatInterface(
     fn=respond,
     chatbot=gr.Chatbot(show_copy_button=True, likeable=True, layout="panel"),
-    description ="# OpenGPT 4o \n ### chat, generate images, perform web searches, and Q&A with images.",
+    description="# OpenGPT 4o \n ### chat, generate images, perform web searches, and Q&A with images.",
     textbox=gr.MultimodalTextbox(),
     multimodal=True,
     concurrency_limit=200,