Update app.py

app.py

@@ -12,14 +12,6 @@ client = OpenAI(
     api_key="hf_XXXXX"
 )
 
-
-# Load model and tokenizer
-'''our_model_path = "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
-
-our_model = AutoModelForCausalLM.from_pretrained(our_model_path, device_map="auto", torch_dtype="auto")
-our_tokenizer = AutoTokenizer.from_pretrained(our_model_path)'''
-
 def format_math(text):
     text = re.sub(r"\[(.*?)\]", r"$$\1$$", text, flags=re.DOTALL)
     text = text.replace(r"\(", "$").replace(r"\)", "$")
@@ -31,68 +23,6 @@ conversations = {}
 def generate_conversation_id():
     return str(uuid.uuid4())[:8]
 
-#@spaces.GPU(duration=60)
-'''def generate_response(user_message, max_tokens, temperature, top_p, history_state):
-    if not user_message.strip():
-        return history_state, history_state
-
-    model = our_model
-    tokenizer = our_tokenizer
-    start_tag = "<|im_start|>"
-    sep_tag = "<|im_sep|>"
-    end_tag = "<|im_end|>"
-
-    system_message = "Your role as an assistant..."
-    prompt = f"{start_tag}system{sep_tag}{system_message}{end_tag}"
-    for message in history_state:
-        if message["role"] == "user":
-            prompt += f"{start_tag}user{sep_tag}{message['content']}{end_tag}"
-        elif message["role"] == "assistant" and message["content"]:
-            prompt += f"{start_tag}assistant{sep_tag}{message['content']}{end_tag}"
-    prompt += f"{start_tag}user{sep_tag}{user_message}{end_tag}{start_tag}assistant{sep_tag}"
-
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-
-    generation_kwargs = {
-        "input_ids": inputs["input_ids"],
-        "attention_mask": inputs["attention_mask"],
-        "max_new_tokens": int(max_tokens),
-        "do_sample": True,
-        "temperature": temperature,
-        "top_k": 50,
-        "top_p": top_p,
-        "repetition_penalty": 1.0,
-        "pad_token_id": tokenizer.eos_token_id,
-        "streamer": streamer,
-    }
-
-    try:
-        thread = Thread(target=model.generate, kwargs=generation_kwargs)
-        thread.start()
-    except Exception:
-        yield history_state + [{"role": "user", "content": user_message}, {"role": "assistant", "content": "⚠️ Generation failed."}], history_state
-        return
-
-    assistant_response = ""
-    new_history = history_state + [
-        {"role": "user", "content": user_message},
-        {"role": "assistant", "content": ""}
-    ]
-
-    try:
-        for new_token in streamer:
-            if "<|end" in new_token:
-                continue
-            cleaned_token = new_token.replace("<|im_start|>", "").replace("<|im_sep|>", "").replace("<|im_end|>", "")
-            assistant_response += cleaned_token
-            new_history[-1]["content"] = assistant_response.strip()
-            yield new_history, new_history
-    except Exception:
-        pass
-
-    yield new_history, new_history
-'''
 
 def generate_response(user_message, max_tokens, temperature, top_p, history_state):
     if not user_message.strip():
@@ -116,6 +46,7 @@ def generate_response(user_message, max_tokens, temperature, top_p, history_stat
             temperature=temperature,
             top_p=top_p,
             stream=True,
+            extra_body={"max_new_tokens": int(max_tokens)}
         )
     except Exception as e:
         print(f"[ERROR] OpenAI API call failed: {e}")