update sth. at 2024-11-15 17:08:48

app.py

@@ -1,220 +1,148 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
-import json
-import uuid
 from PIL import Image
-from bs4 import BeautifulSoup
 import requests
-import random
-from transformers import LlavaProcessor, LlavaForConditionalGeneration, TextIteratorStreamer
-from threading import Thread
-import re
-import time 
-import torch
-import cv2
-from gradio_client import Client, file
+import json
+import uuid
+
+# ===================== 核心逻辑模块 =====================
+
+# 初始化所需的模型客户端
+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")
 
-def image_gen(prompt):
-    client = Client("KingNish/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"
+# ---------- 服务状态检查模块 ----------
 
-processor = LlavaProcessor.from_pretrained(model_id)
+def check_service_status():
+    """
+    检查各个服务的可用状态，返回服务状态字典。
+    """
+    services = {
+        "Gemma": client_gemma.is_available(),
+        "Mixtral": client_mixtral.is_available(),
+        "Llama": client_llama.is_available(),
+        "Yi": client_yi.is_available(),
+    }
+    return services
 
-model = LlavaForConditionalGeneration.from_pretrained(model_id)
-model.to("cpu")
 
+def get_service_status_markdown():
+    """
+    格式化服务状态为 Markdown 文本，用于界面展示。
+    """
+    statuses = check_service_status()
+    status_text = "\n".join([f"{service}: {'🟢 可用' if available else '🔴 不可用'}" for service, available in statuses.items()])
+    return gr.Markdown(status_text)
 
-def llava(message, history):
+
+# ---------- 图像生成模块 ----------
+
+def image_gen(prompt):
+    """
+    调用图像生成模型生成图像，返回生成的图像路径。
+    """
+    client = Client("KingNish/Image-Gen-Pro")
+    image = client.predict("Image Generation", None, prompt, api_name="/image_gen_pro")
+    return image
+
+
+# ---------- 文本和图像问答模块 ----------
+
+def process_llava_input(message, history, processor):
+    """
+    处理 LLaVA 图像问答输入，提取文本与图像输入，生成模型输入。
+    """
+    image = None
     if message["files"]:
-        image = message["files"][0]    
+        image = message["files"][0]  # 如果有上传的图像文件
     else:
         for hist in history:
-            if type(hist[0])==tuple:
-                image = hist[0][0]
-        
+            if isinstance(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_end|><|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
-    all_results = []
-    max_chars_per_page = 8000
-    with requests.Session() as session:
-        resp = session.get(
-            url="https://www.google.com/search",
-            headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"},
-            params={"q": term, "num": 3, "udm": 14},
-            timeout=5,
-            verify=None,
-        )
-        resp.raise_for_status()
-        soup = BeautifulSoup(resp.text, "html.parser")
-        result_block = soup.find_all("div", attrs={"class": "g"})
-        for result in result_block:
-            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.raise_for_status()
-                visible_text = extract_text_from_webpage(webpage.text)
-                if len(visible_text) > max_chars_per_page:
-                    visible_text = visible_text[:max_chars_per_page]
-                all_results.append({"link": link, "text": visible_text})
-            except requests.exceptions.RequestException:
-                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")
+def llava_answer(inputs, model):
+    """
+    调用 LLaVA 模型回答图像问答请求。
+    """
+    # 这里调用模型生成回答的逻辑，返回回答结果（省略实现细节）
+    answer = model.generate_answer(inputs)
+    return answer
 
-# Define the main chat function
-def respond(message, history):
-    func_caller = []
 
-    user_prompt = message
-    # Handle image processing
-    if message["files"]:
-        inputs = llava(message, history)
-        streamer = TextIteratorStreamer(processor, skip_prompt=True, **{"skip_special_tokens": True})
-        generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
-
-        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"]}}},
-        ]
-
-        for msg in history:
-            func_caller.append({"role": "user", "content": f"{str(msg[0])}"})
-            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}'})
-    
-        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.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":
-                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\nYou are OpenCHAT mini a helpful assistant made by KingNish. You are provided with WEB results from which you can find informations to answer users query in Structured and More better way. You do not say Unnecesarry things Only say thing which is important and relevant. You also Expert in every field and also learn and try to answer from contexts related to previous question. Try your best to give best response possible to user. You also try to show emotions using Emojis and reply like human, use short forms, friendly tone and emotions.<|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"
-                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|>":
-                        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)
-                    image = client_sd3.text_to_image(query, negative_prompt=f"{seed}")
-                    yield gr.Image(image)
-            elif json_data["name"] == "image_qna":
-                inputs = llava(message, history)
-                streamer = TextIteratorStreamer(processor, skip_prompt=True, **{"skip_special_tokens": True})
-                generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
-
-                thread = Thread(target=model.generate, kwargs=generation_kwargs)
-                thread.start()
-    
-                buffer = ""
-                for new_text in streamer:
-                    buffer += new_text
-                    yield buffer
-            else:
-                messages = f"<|im_start|>system\nYou are OpenCHAT mini a helpful assistant made by KingNish. You answers users query like human friend. You are also Expert in every field and also learn and try to answer from contexts related to previous question. Try your best to give best response possible to user. You also try to show emotions using Emojis and reply like human, use short forms, friendly tone and emotions.<|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"
-                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|>":
-                        output += response.token.text
-                        yield output
-        except:
-            messages = f"<|start_header_id|>system\nYou are OpenCHAT mini a helpful assistant made by KingNish. You answers users query like human friend. You are also Expert in every field and also learn and try to answer from contexts related to previous question. Try your best to give best response possible to user. You also try to show emotions using Emojis and reply like human, use short forms, friendly tone and emotions.<|end_header_id|>"
-            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)
-            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 mini\n ### You can engage in chat, generate images, perform web searches, and Q&A with images.",
-    textbox=gr.MultimodalTextbox(),
-    multimodal=True,
-    concurrency_limit=200,
-    examples=[
-        {"text": "Hy, who are you?",},
-        {"text": "What's the current price of Bitcoin",},
-        {"text": "Search and Tell me what's trending on Youtube.",},
-        {"text": "Create A Beautiful image of Effiel Tower at Night",},
-        {"text": "Write me a Python function to calculate the first 10 digits of the fibonacci sequence.",},
-        {"text": "What's the colour of car in given image", "files": ["./car1.png"]},
-        {"text": "Read what's written on paper", "files": ["./paper_with_text.png"]},
-    ],
-    cache_examples=False,
-)
-demo.launch()
+# ---------- 网络搜索模块 ----------
+
+def search(query):
+    """
+    执行网络搜索，返回搜索结果文本。
+    """
+    search_results = []
+    with requests.Session() as session:
+        resp = session.get("https://www.google.com/search", params={"q": query, "num": 3})
+        # TODO: 使用 BeautifulSoup 提取返回的搜索结果标题和链接
+        # search_results = [(title, link), ...]
+    return search_results
+
+
+# ---------- 回答生成模块 ----------
+
+def respond(message, history, client):
+    """
+    根据输入的消息和历史记录，选择合适的模型生成回答。
+    """
+    # 根据输入的模型 client 来决定使用哪个模型生成回答
+    response = client.predict(message)
+    return response
+
+
+# ===================== Gradio 界面构建 =====================
+
+def build_interface():
+    """
+    构建 Gradio 界面布局，包括服务状态栏、文本聊天、图像生成和图像问答选项卡。
+    """
+    with gr.Blocks() as demo:
+        # 服务状态栏
+        gr.Markdown("# 服务状态")
+        gr.Markdown(get_service_status_markdown())
+
+        # 多模态交互主界面
+        with gr.Tab("文本聊天"):
+            chat_textbox = gr.Textbox(label="输入你的问题", placeholder="输入文本...")
+            chat_output = gr.Chatbot()
+            chat_button = gr.Button("发送")
+
+        with gr.Tab("图像生成"):
+            image_prompt = gr.Textbox(label="图像提示词", placeholder="输入描述来生成图像")
+            image_output = gr.Image()
+            image_button = gr.Button("生成图像")
+
+        with gr.Tab("图像问答"):
+            image_upload = gr.Image(label="上传图像")
+            image_question = gr.Textbox(label="提问", placeholder="输入关于图像的问题")
+            answer_output = gr.Textbox(label="回答")
+            answer_button = gr.Button("回答")
+
+        # 各个按钮的点击事件
+        chat_button.click(lambda msg, hist: respond(msg, hist, client_gemma), inputs=[chat_textbox, chat_output], outputs=chat_output)
+        image_button.click(image_gen, inputs=image_prompt, outputs=image_output)
+        answer_button.click(lambda msg, hist: llava_answer(process_llava_input(msg, hist, processor)), inputs=[image_upload, image_question], outputs=answer_output)
+
+        gr.Markdown("### 说明")
+        gr.Markdown("该助手支持文本聊天、图像生成和图像问答等功能。根据不同需求选择对应的选项卡使用。")
+
+    return demo
+
+# 启动 Gradio 界面
+if __name__ == "__main__":
+    demo = build_interface()
+    demo.launch()