Create app.py

app.py

@@ -0,0 +1,262 @@
+import os
+import time
+import spaces
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import gradio as gr
+from threading import Thread
+from huggingface_hub import login
+from icrawler.builtin import BingImageCrawler
+import os
+
+MODEL_LIST = ["mistralai/Mistral-Nemo-Instruct-2407"]
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+login(token=HF_TOKEN)
+#MODEL = os.environ.get("MODEL_ID")
+MODEL = "mistralai/Mistral-Nemo-Instruct-2407"
+
+TITLE = "<h1><center>Mistral-Nemo</center></h1>"
+
+PLACEHOLDER = """
+<center>
+<p>The Mistral-Nemo is a pretrained generative text model of 12B parameters trained jointly by Mistral AI and NVIDIA.</p>
+</center>
+"""
+
+
+CSS = """
+.duplicate-button {
+    margin: auto !important;
+    color: white !important;
+    background: black !important;
+    border-radius: 100vh !important;
+}
+h3 {
+    text-align: center;
+}
+
+#output_video {
+    display: block;
+    margin-left: auto!important;
+    margin-right: auto !important;
+    width: 20vw !important;
+}
+
+footer{visibility: hidden}
+"""
+
+device = "cuda" # for GPU usage or "cpu" for CPU usage
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    ignore_mismatched_sizes=True)
+
+@spaces.GPU()
+def get_response(conversation):
+    temperature  = 0.3
+    max_new_tokens = 512
+    top_p = 1.0
+    top_k = 20 
+    penalty = 1.2
+
+    input_text=tokenizer.apply_chat_template(conversation, tokenize=False)
+    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
+    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+    
+    generate_kwargs = dict(
+        input_ids=inputs, 
+        max_new_tokens = max_new_tokens,
+        do_sample = False if temperature == 0 else True,
+        top_p = top_p,
+        top_k = top_k,
+        temperature = temperature,
+        streamer=streamer,
+        repetition_penalty=penalty,
+        pad_token_id = 10,
+    )
+
+    with torch.no_grad():
+        thread = Thread(target=model.generate, kwargs=generate_kwargs)
+        thread.start()
+        
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text
+        #yield buffer
+
+    return buffer
+    
+
+@spaces.GPU()
+def stream_chat(history, character_a, character_b):
+
+    # From A to B
+    if(history==[]):
+        conversation_a = [{"role": "system", "content": "You should respond like " + character_b + ". You should have a meaningful conversation. Don't repeat yourself. You should only output your response. You don't need to put quotes around what your saying. You don't need to put your name at the beginning of your response."}]
+        for answer, prompt in history:
+            conversation_a.extend([
+                {"role": "user", "content": prompt}, 
+                {"role": "assistant", "content": answer},
+            ])
+        conversation_a.append({"role": "user", "content": "You are having a conversation with " +  character_a + ". Introduce yourself."})
+        response_b = get_response(conversation_a)
+        print('response_b', response_b)
+    else:
+        conversation_a = [{"role": "system", "content": "You should respond like " + character_b + ". You should have a meaningful conversation. Don't repeat yourself. You should only output your response. You don't need to put quotes around what your saying. You don't need to put your name at the beginning of your response."}]
+        for answer, prompt in history:
+            conversation_a.extend([
+                {"role": "user", "content": prompt}, 
+                {"role": "assistant", "content": answer},
+            ])
+        conversation_a.append({"role": "user", "content": history[-1][0] })
+        response_b = get_response(conversation_a)
+        print('response_b', response_b)
+        
+
+    # From B to A
+    conversation_b = [{"role": "system", "content": "You should respond like " + character_a + ". You should have a meaningful conversation. Don't repeat yourself. You should only output your response. You don't need to put quotes around what your saying. You don't need to put your name at the beginning of your response."}]
+    for prompt, answer in history:
+        conversation_b.extend([
+            {"role": "user", "content": prompt}, 
+            {"role": "assistant", "content": answer},
+        ])
+    conversation_b.append({"role": "user", "content": response_b })
+    response_a = get_response(conversation_b)
+    print('response_a', response_a)
+
+    # Append responses to history
+    history.append((response_b , response_a ))
+    print('history', history)
+    
+    return history
+
+
+def get_img(keyword):
+    path = './' + keyword
+    os.makedirs(path, exist_ok=True)
+    bing_crawler = BingImageCrawler(storage={'root_dir': path})
+    bing_crawler.crawl(keyword=keyword, max_num=1)
+
+    # Look for image files in the folder
+    for file_name in os.listdir(path):
+        if file_name.lower().endswith(('.png', '.PNG', '.JPG', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff')):
+            return os.path.join(path, file_name)
+    
+    # If no image is found
+    return None
+
+def set_characters(a, b):
+    img_a = get_img(a)
+    img_b =  get_img(b)
+    return img_a, img_b, gr.update(avatar_images=(img_b, img_a))
+            
+chatbot = gr.Chatbot(height=600,show_label=False)
+
+theme = gr.themes.Base(
+).set(
+    body_background_fill="#e1fceb",
+    color_accent_soft="#ffffff",
+    border_color_accent="#e1fceb",
+    border_color_primary="#e1fceb",
+    background_fill_secondary="#e1fceb",
+    button_secondary_background_fill="#ffffff",
+    button_primary_background_fill="#ffffff",
+    button_primary_text_color="#1f2937",
+    input_background_fill="#f8f8f8",
+)
+
+with gr.Blocks(css=CSS, theme=theme) as demo:
+    gr.HTML("""
+
+        <center> <h1> Bot vs Bot </h1> </center>
+    
+        <center> by <a href="https://www.tonyassi.com/">Tony Assi</a> </center>
+        
+        <center> <h3> Pick two icons and watch them have a conversation </h3> </center>
+    """)
+    with gr.Row():
+        character_a = gr.Textbox(label='Character A', info='Choose a person', placeholder='Socrates, Edgar Allen Poe, George Washington')
+        character_b = gr.Textbox(label='Character B', info='Choose a person', placeholder='Madonna, Paris Hilton, Liza Minnelli')
+    character_button = gr.Button('Initiate Characters')
+    #characters_html = gr.HTML()
+    with gr.Row():
+        image_a = gr.Image(show_label=False, interactive=False, show_fullscreen_button=False, show_share_button=False, show_download_button=False)
+        gr.Markdown(' ')
+        image_b = gr.Image(show_label=False, interactive=False, show_fullscreen_button=False, show_share_button=False, show_download_button=False)
+    
+
+    chat = gr.Chatbot(show_label=False)
+    submit_button = gr.Button()
+
+    character_button.click(set_characters, inputs=[character_a, character_b], outputs=[image_a, image_b, chat])
+    submit_button.click(stream_chat, inputs=[chat, character_a, character_b], outputs=[chat])
+    
+    
+    #gr.HTML(TITLE)
+    #gr.DuplicateButton(value="Duplicate Space for private use", elem_classes="duplicate-button")
+    #video = gr.HTML("""
+    #  <video src="https://huggingface.co/spaces/ahjvdjf33/moods/resolve/main/Paris-Moods-2/passive3.mp4"
+    #         playsinline autoplay loop muted></video>
+    #""", elem_id="output_video")
+    """
+    gr.ChatInterface(
+        fn=stream_chat,
+        chatbot=chatbot,
+        textbox=gr.Textbox(visible=False),
+        fill_height=True,
+        retry_btn=None,
+        undo_btn=None,
+        clear_btn=None,
+        additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False, visible=False),
+        additional_inputs=[
+            gr.Slider(
+                minimum=0,
+                maximum=1,
+                step=0.1,
+                value=0.3,
+                label="Temperature",
+                render=False,
+            ),
+            gr.Slider(
+                minimum=128,
+                maximum=8192,
+                step=1,
+                value=1024,
+                label="Max new tokens",
+                render=False,
+            ),
+            gr.Slider(
+                minimum=0.0,
+                maximum=1.0,
+                step=0.1,
+                value=1.0,
+                label="top_p",
+                render=False,
+            ),
+            gr.Slider(
+                minimum=1,
+                maximum=20,
+                step=1,
+                value=20,
+                label="top_k",
+                render=False,
+            ),
+            gr.Slider(
+                minimum=0.0,
+                maximum=2.0,
+                step=0.1,
+                value=1.2,
+                label="Repetition penalty",
+                render=False,
+            ),
+        ],
+    )
+    """
+
+
+
+if __name__ == "__main__":
+    demo.launch()