Update app.py

app.py

@@ -1,25 +1,14 @@
-import os
-
-os.system("pip uninstall -y gradio")
-os.system("pip install gradio==4.44.1") 
-os.system("pip install langfuse")
-
-from threading import Thread
-from typing import Iterator
-
 import gradio as gr
-from langfuse import Langfuse
-from langfuse.decorators import observe
-import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
+from typing import Iterator
 import time
-
-#from utils import load_list_from_json
+import os
 
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
-MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
+MAX_INPUT_TOKEN_LENGTH = 4096
 
 DESCRIPTION = """\
 # ISANG-1.0-8B Chat
@@ -47,19 +36,16 @@ pre, code {
 }
 """
 
-system_prompt = str(os.getenv("SYSTEM_PROMPT"))
+system_prompt = "You are a helpful assistant."
 
-secret_key = str(os.getenv("LANGFUSE_SECRET_KEY"))
-public_key = str(os.getenv("LANGFUSE_PUBLIC_KEY"))
-host = str(os.getenv("LANGFUSE_HOST"))
-
-langfuse = Langfuse(
-  secret_key=secret_key,
-  public_key=public_key,
-  host=host
+# 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)
 
-#REJECTED_VOCAB = load_list_from_json("rejected_vocab_extended.json")
+generation_speed = 0
 
 def execution_time_calculator(start_time, log=True):
     delta = time.time() - start_time
@@ -70,30 +56,6 @@ def execution_time_calculator(start_time, log=True):
 def token_per_second_calculator(tokens_count, time_delta):
     return tokens_count / time_delta
 
-if not torch.cuda.is_available():
-    DESCRIPTION = "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
-
-if torch.cuda.is_available():
-    # Use your own model here:
-    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 get_generation_speed():
-    global generation_speed
-    return generation_speed
-
-@observe()
-def log_to_langfuse(message, chat_history, max_new_tokens, temperature, top_p, top_k, 
-                    repetition_penalty, do_sample, generation_speed, model_outputs):
-    print(f"generation_speed: {generation_speed}")
-    return "".join(model_outputs)
-
-@spaces.GPU
 def generate(
     message: str,
     chat_history: list[tuple[str, str]],
@@ -104,24 +66,18 @@ def generate(
     repetition_penalty: float = 1.2,
     do_sample: bool = True,
 ) -> Iterator[str]:
-    global generation_speed
-    global system_prompt
-
     conversation = []
-    if system_prompt:
-        conversation.append({"role": "system", "content": system_prompt})
-    for user, assistant in chat_history:
-        conversation.extend([
-            {"role": "user", "content": user}, 
-            {"role": "assistant", "content": assistant}
-        ])
+    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.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)
+    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
@@ -136,7 +92,6 @@ def generate(
         temperature=temperature,
         num_beams=1,
         repetition_penalty=repetition_penalty,
-        #bad_words_ids=REJECTED_VOCAB,
     )
 
     start_time = time.time()
@@ -148,36 +103,23 @@ def generate(
     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)
 
-    log_function = log_to_langfuse(
-        message=message,
-        chat_history=chat_history,
-        max_new_tokens=max_new_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        top_k=top_k,
-        repetition_penalty=repetition_penalty,
-        do_sample=do_sample,
-        generation_speed=generation_speed,
-        model_outputs=outputs,
-    )
-
-chatbot = gr.Chatbot(placeholder=PLACEHOLDER, scale=1, show_copy_button=True, height="68%", rtl=True)
-chat_input = gr.Textbox(show_label=False, lines=2, rtl=True, placeholder="ورودی", show_copy_button=True, scale=4)
-submit_btn = gr.Button(variant="primary", value="ارسال", size="sm", scale=1, elem_classes=["_button"])
+# 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="ورودی‌های اضافی", open=False),
+    additional_inputs_accordion=gr.Accordion(label="Additional Inputs", open=False),
     additional_inputs=[
         gr.Slider(
-            label="حداکثر تعداد توکن ها",
+            label="Max New Tokens",
             minimum=1,
             maximum=MAX_MAX_NEW_TOKENS,
             step=1,
@@ -188,7 +130,7 @@ chat_interface = gr.ChatInterface(
             minimum=0.01,
             maximum=4.0,
             step=0.01,
-            value=0.5,
+            value=0.6,
         ),
         gr.Slider(
             label="Top-p",
@@ -202,28 +144,28 @@ chat_interface = gr.ChatInterface(
             minimum=1,
             maximum=1000,
             step=1,
-            value=20,
+            value=50,
         ),
         gr.Slider(
-            label="جریمه تکرار",
+            label="Repetition Penalty",
             minimum=1.0,
             maximum=2.0,
             step=0.05,
             value=1.2,
         ),
         gr.Dropdown(
-            label="نمونه‌گیری",
+            label="Sampling",
             choices=[False, True],
             value=True
         )
     ],
-    stop_btn="توقف",
+    stop_btn="Stop",
     chatbot=chatbot,
     textbox=chat_input,
     submit_btn=submit_btn,
-    retry_btn="🔄 تلاش مجدد",
-    undo_btn="↩️ بازگشت",
-    clear_btn="🗑️ پاک کردن",
+    retry_btn="Retry",
+    undo_btn="Undo",
+    clear_btn="Clear",
     title="ISANG AI"
 )