Update app.py

app.py

@@ -1,178 +1,42 @@
 import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
-from typing import Iterator
-import time
-import os
-
-MAX_MAX_NEW_TOKENS = 2048
-DEFAULT_MAX_NEW_TOKENS = 1024
-MAX_INPUT_TOKEN_LENGTH = 4096
-
-DESCRIPTION = """\
-# ISANG-1.0-8B Chat
-"""
-
-PLACEHOLDER = """
-<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
-   <img src="https://avatars.githubusercontent.com/u/39557177?v=4" 
-        style="width: 80%; max-width: 550px; height: auto; opacity: 0.80;"> 
-   <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">ISANG-1.0-8B</h1>
-</div>
-"""
-
-custom_css = """
-@import url('https://fonts.googleapis.com/css2?family=Vazirmatn&display=swap');
-body, .gradio-container, .gr-button, .gr-input, .gr-slider, .gr-dropdown, .gr-markdown {
-    font-family: 'Vazirmatn', sans-serif !important;
-}
-._button {
-    font-size: 20px;
-}
-pre, code {
-    direction: ltr !important;
-    unicode-bidi: plaintext !important;
-}
-"""
-
-system_prompt = "You are a helpful assistant."
-
-# Load the model
-model_id = "hosseinhimself/ISANG-1.0-8B"
-model = AutoModelForCausalLM.from_pretrained(
-    model_id, device_map="auto", torch_dtype=torch.bfloat16
-)
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-
-generation_speed = 0
-
-def execution_time_calculator(start_time, log=True):
-    delta = time.time() - start_time
-    if log:
-        print("--- %s seconds ---" % (delta))
-    return delta
-
-def token_per_second_calculator(tokens_count, time_delta):
-    return tokens_count / time_delta
-
-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,
-    do_sample: bool = True,
-) -> Iterator[str]:
-    conversation = []
-    conversation.append({"role": "system", "content": system_prompt})
-
-    # Add previous conversation to history (send the last two exchanges as context)
-    for user, assistant in chat_history[-2:]:
-        conversation.append({"role": "user", "content": user})
-        conversation.append({"role": "assistant", "content": assistant})
-
-    conversation.append({"role": "user", "content": message})
-
-    input_ids = tokenizer(conversation, return_tensors="pt", padding=True, truncation=True)
-    input_ids = 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=do_sample,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        num_beams=1,
-        repetition_penalty=repetition_penalty,
-    )
-
-    start_time = time.time()
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    outputs = []
-    sum_tokens = 0
-    for text in streamer:
-        num_tokens = len(tokenizer.tokenize(text))
-        sum_tokens += num_tokens
-        outputs.append(text)
-        yield "".join(outputs)
-
-    time_delta = execution_time_calculator(start_time, log=False)
-    generation_speed = token_per_second_calculator(sum_tokens, time_delta)
-
-# Define Gradio interface components
-chatbot = gr.Chatbot(placeholder=PLACEHOLDER, scale=1, show_copy_button=True, height="68%")
-chat_input = gr.Textbox(show_label=False, lines=2, placeholder="Enter your message", show_copy_button=True, scale=4)
-submit_btn = gr.Button(variant="primary", value="Submit", size="sm", scale=1, elem_classes=["_button"])
-
-chat_interface = gr.ChatInterface(
-    fn=generate,
-    additional_inputs_accordion=gr.Accordion(label="Additional Inputs", open=False),
-    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.01,
-            maximum=4.0,
-            step=0.01,
-            value=0.6,
-        ),
-        gr.Slider(
-            label="Top-p",
-            minimum=0.05,
-            maximum=1.0,
-            step=0.01,
-            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,
-        ),
-        gr.Dropdown(
-            label="Sampling",
-            choices=[False, True],
-            value=True
-        )
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# Load model and tokenizer
+model_name = "hosseinhimself/ISANG-v1.0-8B"
+model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Set hyperparameters for inference
+def generate_response(input_text, max_tokens=1024, temperature=0.7, history=[]):
+    # Prepare the inputs for the model
+    prompt = f"Chat History: {history[-2:]}\nUser: {input_text}\nAI:"
+    
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    output = model.generate(**inputs, max_new_tokens=max_tokens, temperature=temperature)
+    
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    
+    # Update history with the new conversation
+    history.append(f"User: {input_text}")
+    history.append(f"AI: {response}")
+    
+    return response, history
+
+# Gradio interface
+iface = gr.Interface(
+    fn=generate_response,
+    inputs=[
+        gr.Textbox(label="Your Message", placeholder="Type your message here..."),
+        gr.Slider(minimum=1, maximum=2048, value=1024, step=1, label="Max Tokens"),
+        gr.Slider(minimum=0.0, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
+        gr.State(value=[])
     ],
-    stop_btn="Stop",
-    chatbot=chatbot,
-    textbox=chat_input,
-    submit_btn=submit_btn,
-    retry_btn="Retry",
-    undo_btn="Undo",
-    clear_btn="Clear",
-    title="ISANG AI"
+    outputs=[gr.Textbox(label="AI Response"), gr.State()],
+    title="ISANG Chatbot",
+    description="A chatbot powered by ISANG-v1.0-8B model. Chat with me!",
+    theme="huggingface",
+    live=True
 )
 
-# --- Use Gradio's built-in theming system with a purple accent ---
-with gr.Blocks(css=custom_css, fill_height=False, theme=gr.themes.Default(primary_hue="purple")) as demo:
-    gr.Markdown(DESCRIPTION)
-    chat_interface.render()
-
-if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+# Launch the interface
+iface.launch()