DEV: Use mistral template

app.py

@@ -1,4 +1,5 @@
-# Importing libraries
+#!/usr/bin/env python#
+
 from llama_cpp import Llama
 from time import time
 import gradio as gr
@@ -20,63 +21,134 @@ It uses quantized models in gguf-Format and llama.cpp to run them.
 <br>''' + f"DEBUG: Memory used: {psutil.virtual_memory()[2]}<br>" + '''
 Powered by ...'''
 
-# Loading prompt
-prompt = ""
-system_message = ""
-
-def generate_answer(request: str, max_tokens: int = 256, custom_prompt: str = None):
-    t0 = time()
-    logs = f"Request: {request}\nMax tokens: {max_tokens}\nCustom prompt: {custom_prompt}\n"
-    try:
-        maxTokens = max_tokens if 16 <= max_tokens <= 256 else 64
-        userPrompt = prompt.replace("{prompt}", request)
-        userPrompt = userPrompt.replace(
-            "{system_message}",
-            custom_prompt if isinstance(custom_prompt, str) and len(custom_prompt.strip()) > 1 and custom_prompt.strip() not in ['', None, ' '] else system_message
-        )
-        logs += f"\nFinal prompt: {userPrompt}\n"
-    except:
-        return "Not enough data! Check that you passed all needed data.", logs
-    
-    try:
-        # this shitty fix will be until i willnt figure out why sometimes there is empty output
-        counter = 1
-        while counter <= 3:
-            logs += f"Attempt {counter} to generate answer...\n"
-            output = llm(userPrompt, max_tokens=maxTokens, stop=["<|im_end|>"], echo=False)
-            text = output["choices"][0]["text"]
-            if len(text.strip()) > 1 and text.strip() not in ['', None, ' ']:
-                break
-            counter += 1
-        logs += f"Final attempt: {counter}\n"
-        if len(text.strip()) <= 1 or text.strip() in ['', None, ' ']:
-            logs += f"Generated and aborted: {text}"
-            text = "Sorry, but something went wrong while generating answer. Try again or fix code. If you are maintainer of this space, look into logs."
-        
-        logs += f"\nFinal: '''{text}'''"
-        logs += f"\n\nTime spent: {time()-t0}"
-        return text, logs
-    except Exception as e:
-        logs += str(e)
-        logs += f"\n\nTime spent: {time()-t0}"
-        return "Oops! Internal server error. Check the logs of space/instance.", logs
-
-print("! LOAD GRADIO INTERFACE !")
-demo = gr.Interface(
-    fn=generate_answer,
-    inputs=[
-        gr.components.Textbox(label="Input"),
-        gr.components.Number(value=256),
-        gr.components.Textbox(label="Custom system prompt"),
+
+
+import os
+from threading import Thread
+from typing import Iterator
+
+import gradio as gr
+import spaces
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+
+DESCRIPTION = "# Mistral-7B"
+
+if torch.cuda.is_available():
+    DESCRIPTION += "\n<p>This space is optimized for CPU only. Use another one if you want to go fast and use GPU. </p>"
+
+MAX_MAX_NEW_TOKENS = 2048
+DEFAULT_MAX_NEW_TOKENS = 1024
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
+
+
+#download model here
+# check localstorage, if no there, load, else use existing.
+
+if torch.cuda.is_available():
+    model_id = "mistralai/Mistral-7B-Instruct-v0.1"
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+
+def generate(
+    message: str,
+    chat_history: list[tuple[str, str]],
+    max_new_tokens: int = 1024,
+    temperature: float = 0.6,
+    top_p: float = 0.9,
+    top_k: int = 50,
+    repetition_penalty: float = 1.2,
+) -> Iterator[str]:
+    conversation = []
+    for user, assistant in chat_history:
+        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
+    conversation.append({"role": "user", "content": message})
+
+    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
+    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
+        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
+        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
+    input_ids = input_ids.to(model.device)
+
+    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+    generate_kwargs = dict(
+        {"input_ids": input_ids},
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        top_p=top_p,
+        top_k=top_k,
+        temperature=temperature,
+        num_beams=1,
+        repetition_penalty=repetition_penalty,
+    )
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+        yield "".join(outputs)
+
+
+chat_interface = gr.ChatInterface(
+    fn=generate,
+    additional_inputs=[
+        gr.Slider(
+            label="Max new tokens",
+            minimum=1,
+            maximum=MAX_MAX_NEW_TOKENS,
+            step=1,
+            value=DEFAULT_MAX_NEW_TOKENS,
+        ),
+        gr.Slider(
+            label="Temperature",
+            minimum=0.1,
+            maximum=4.0,
+            step=0.1,
+            value=0.6,
+        ),
+        gr.Slider(
+            label="Top-p (nucleus sampling)",
+            minimum=0.05,
+            maximum=1.0,
+            step=0.05,
+            value=0.9,
+        ),
+        gr.Slider(
+            label="Top-k",
+            minimum=1,
+            maximum=1000,
+            step=1,
+            value=50,
+        ),
+        gr.Slider(
+            label="Repetition penalty",
+            minimum=1.0,
+            maximum=2.0,
+            step=0.05,
+            value=1.2,
+        ),
     ],
-    outputs=[
-        gr.components.Textbox(label="Output"),
-        gr.components.Textbox(label="Logs")
+    stop_btn=None,
+    examples=[
+        ["Hello there! How are you doing?"],
+        ["Can you explain briefly to me what is the Python programming language?"],
+        ["Explain the plot of Cinderella in a sentence."],
+        ["How many hours does it take a man to eat a Helicopter?"],
+        ["Write a 100-word article on 'Benefits of Open-Source in AI research'"],
     ],
-    title=title,
-    description=desc,
-    allow_flagging='never'
 )
-demo.queue()
-print("! LAUNCHING GRADIO !")
-demo.launch(server_name="0.0.0.0")
+
+with gr.Blocks(css="style.css") as demo:
+    gr.Markdown(DESCRIPTION)
+    gr.DuplicateButton(
+        value="Duplicate Space for private use",
+        elem_id="duplicate-button",
+        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
+    )
+    chat_interface.render()
+
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch()