|
Here is an improved and fixed version of the `app.py` script with better structure, error handling, and documentation: |
|
|
|
```python |
|
import uuid |
|
import gradio as gr |
|
import re |
|
from diffusers.utils import load_image |
|
import requests |
|
from awesome_chat import chat_huggingface |
|
import os |
|
|
|
|
|
os.makedirs("public/images", exist_ok=True) |
|
os.makedirs("public/audios", exist_ok=True) |
|
os.makedirs("public/videos", exist_ok=True) |
|
|
|
|
|
HUGGINGFACE_TOKEN = os.environ.get("HUGGINGFACE_TOKEN") |
|
OPENAI_KEY = os.environ.get("OPENAI_KEY") |
|
|
|
class Client: |
|
def __init__(self) -> None: |
|
self.OPENAI_KEY = OPENAI_KEY |
|
self.HUGGINGFACE_TOKEN = HUGGINGFACE_TOKEN |
|
self.all_messages = [] |
|
|
|
def set_key(self, openai_key): |
|
self.OPENAI_KEY = openai_key |
|
return self.OPENAI_KEY |
|
|
|
def set_token(self, huggingface_token): |
|
self.HUGGINGFACE_TOKEN = huggingface_token |
|
return self.HUGGINGFACE_TOKEN |
|
|
|
def add_message(self, content, role): |
|
message = {"role": role, "content": content} |
|
self.all_messages.append(message) |
|
|
|
def extract_medias(self, message): |
|
url_pattern = re.compile(r"(http(s?):|\/)?([\.\/_\w:-])*?\.(jpg|jpeg|tiff|gif|png|flac|wav|mp4)") |
|
urls = [match.group(0) for match in url_pattern.finditer(message)] |
|
|
|
image_pattern = re.compile(r"(http(s?):|\/)?([\.\/_\w:-])*?\.(jpg|jpeg|tiff|gif|png)") |
|
audio_pattern = re.compile(r"(http(s?):|\/)?([\.\/_\w:-])*?\.(flac|wav)") |
|
video_pattern = re.compile(r"(http(s?):|\/)?([\.\/_\w:-])*?\.(mp4)") |
|
|
|
image_urls = [match.group(0) for match in image_pattern.finditer(message)] |
|
audio_urls = [match.group(0) for match in audio_pattern.finditer(message)] |
|
video_urls = [match.group(0) for match in video_pattern.finditer(message)] |
|
|
|
return urls, image_urls, audio_urls, video_urls |
|
|
|
def add_text(self, messages, message): |
|
if not self.OPENAI_KEY or not self.OPENAI_KEY.startswith("sk-") or not self.HUGGINGFACE_TOKEN or not self.HUGGINGFACE_TOKEN.startswith("hf_"): |
|
return messages, "Please set your OpenAI API key and Hugging Face token first!" |
|
|
|
self.add_message(message, "user") |
|
messages.append((message, None)) |
|
urls, image_urls, audio_urls, video_urls = self.extract_medias(message) |
|
|
|
for image_url in image_urls: |
|
image_url = "public/" + image_url if not image_url.startswith("http") else image_url |
|
image = load_image(image_url) |
|
name = f"public/images/{str(uuid.uuid4())[:4]}.jpg" |
|
image.save(name) |
|
messages.append(((f"{name}",), None)) |
|
|
|
for audio_url in audio_urls: |
|
audio_url = "public/" + audio_url if not audio_url.startswith("http") else audio_url |
|
ext = audio_url.split(".")[-1] |
|
name = f"public/audios/{str(uuid.uuid4())[:4]}.{ext}" |
|
response = requests.get(audio_url) |
|
with open(name, "wb") as f: |
|
f.write(response.content) |
|
messages.append(((f"{name}",), None)) |
|
|
|
for video_url in video_urls: |
|
video_url = "public/" + video_url if not video_url.startswith("http") else video_url |
|
ext = video_url.split(".")[-1] |
|
name = f"public/videos/{str(uuid.uuid4())[:4]}.{ext}" |
|
response = requests.get(video_url) |
|
with open(name, "wb") as f: |
|
f.write(response.content) |
|
messages.append(((f"{name}",), None)) |
|
|
|
return messages, "" |
|
|
|
def bot(self, messages): |
|
if not self.OPENAI_KEY or not self.OPENAI_KEY.startswith("sk-") or not self.HUGGINGFACE_TOKEN or not self.HUGGINGFACE_TOKEN.startswith("hf_"): |
|
return messages, {} |
|
|
|
message, results = chat_huggingface(self.all_messages, self.OPENAI_KEY, self.HUGGINGFACE_TOKEN) |
|
_, image_urls, audio_urls, video_urls = self.extract_medias(message) |
|
self.add_message(message, "assistant") |
|
messages[-1][1] = message |
|
|
|
for image_url in image_urls: |
|
if not image_url.startswith("http"): |
|
image_url = "public/" + image_url.replace("public/", "") |
|
messages.append((None, f |