Create app.py

app.py

@@ -0,0 +1,136 @@
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+model_path = "BytArch/source-mini"
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForCausalLM.from_pretrained(model_path)
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+def generate_response(
+    prompt,
+    system_message,
+    conversation_history=None,
+    max_tokens=75,
+    temperature=0.78,
+    top_p=0.85,
+    repetition_penalty=1.031,
+    top_k=55,
+):
+    context = ""
+    if conversation_history:
+        recent = conversation_history[-30:] if len(conversation_history) > 30 else conversation_history
+        is_first_message = False
+        for i, message in enumerate(recent):
+            if i == 0:
+                is_first_message = True
+                context += (
+                    f"<|start|>User:<|message|>{system_message}<|end|>\n"
+                    f"<|start|>Assistant:<|message|>Hello, nice to meet you!<|end|>\n"
+                )
+            if message["role"] == "user":
+                context += f"<|start|>User:<|message|>{message['content']}<|end|>\n"
+            else:
+                context += f"<|start|>Assistant:<|message|>{message['content']}<|end|>\n"
+
+    formatted_input = (
+        f"{context}<|start|>User:<|message|>{prompt}<|end|>\n<|start|>Assistant:<|message|>"
+    )
+
+    inputs = tokenizer(
+        formatted_input,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=512,
+    )
+
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs.input_ids,
+            attention_mask=inputs.attention_mask,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            do_sample=True,
+            pad_token_id=tokenizer.pad_token_id,
+            repetition_penalty=repetition_penalty,
+            eos_token_id=tokenizer.encode("<|end|>", add_special_tokens=False)[0],
+        )
+
+    new_tokens = outputs[0][inputs.input_ids.shape[-1]:]
+    response = tokenizer.decode(new_tokens, skip_special_tokens=False)
+
+    return response.strip()
+
+def respond(
+    message,
+    history: list[dict[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+    repetition_penalty,
+    top_k,
+):
+    conversation_history = history
+    response = generate_response(
+        message,
+        system_message,
+        conversation_history,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        top_k=top_k,
+    )
+
+    if "<|end|>" in response:
+        response = response.split("<|end|>")[0]
+
+    return response.strip()
+
+chatbot = gr.ChatInterface(
+    respond,
+    type="messages",
+    title="Chat with source-mini",
+    description="Chat with BytArch/source-mini",
+    additional_inputs=[
+        gr.Textbox(
+            value="You are source-mini, a helpful medical/nursing assistant chatbot.",
+            label="System message",
+        ),
+        gr.Slider(minimum=10, maximum=150, value=75, step=5, label="Max new tokens"),
+        gr.Slider(minimum=0.01, maximum=1.2, value=0.7, step=0.01, label="Temperature"),
+        gr.Slider(
+            minimum=0.01,
+            maximum=1.0,
+            value=0.85,
+            step=0.01,
+            label="Top-p (nucleus sampling)",
+        ),
+        gr.Slider(
+            minimum=1.0,
+            maximum=1.5,
+            value=1.031,
+            step=0.001,
+            label="Repetition penalty",
+        ),
+        gr.Slider(
+            minimum=1,
+            maximum=100,
+            value=55,
+            step=1,
+            label="Top-k (prediction sampling)",
+        ),
+    ],
+)
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    chatbot.render()
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False, debug=True)