Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import json
 import subprocess
 from threading import Thread
@@ -10,10 +9,10 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-MODEL_ID = "nikravan/Marco_o1_q4"
+MODEL_ID = "AtAndDev/marco-qwq-7B"
 CHAT_TEMPLATE = "ChatML"
 MODEL_NAME = MODEL_ID.split("/")[-1]
-CONTEXT_LENGTH = 16000
+CONTEXT_LENGTH = 32000
 
 # Estableciendo valores directamente para las variables
 COLOR = "blue"  # Color predeterminado de la interfaz
@@ -75,7 +74,7 @@ def predict(message, history, system_prompt, temperature, max_new_tokens, top_k,
     else:
         raise Exception("Incorrect chat template, select 'Auto', 'ChatML' or 'Mistral Instruct'")
     print(instruction)
-    
+
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     enc = tokenizer(instruction, return_tensors="pt", padding=True, truncation=True)
     input_ids, attention_mask = enc.input_ids, enc.attention_mask
@@ -104,14 +103,13 @@ def predict(message, history, system_prompt, temperature, max_new_tokens, top_k,
             break
         yield "".join(outputs)
 
-
 # Load model
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 quantization_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_compute_dtype=torch.bfloat16
 )
-tokenizer = AutoTokenizer.from_pretrained('AIDC-AI/Marco-o1')
+tokenizer = AutoTokenizer.from_pretrained('AtAndDev/marco-qwq-7B')
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     device_map="auto",
@@ -124,12 +122,10 @@ gr.ChatInterface(
     predict,
     title=EMOJI + " " + MODEL_NAME,
     description=DESCRIPTION,
-    
 
-     
     additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False),
     additional_inputs=[
-        gr.Textbox("You are a code assistant.", label="System prompt"),
+        gr.Textbox("You are a helpful assistant.", label="System prompt"),
         gr.Slider(0, 1, 0.3, label="Temperature"),
         gr.Slider(128, 4096, 1024, label="Max new tokens"),
         gr.Slider(1, 80, 40, label="Top K sampling"),
@@ -137,5 +133,4 @@ gr.ChatInterface(
         gr.Slider(0, 1, 0.95, label="Top P sampling"),
     ],
     theme=gr.themes.Soft(primary_hue=COLOR),
-).queue().launch()
-
+).queue().launch()