Update app.py

app.py

@@ -4,13 +4,11 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 MODEL_ID = "vietrix/viena-60m"
 
-# ==== LOAD MODEL 1 LẦN LÚC START SPACE ====
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto",  # cho nó tự nhét lên GPU nếu có
+    torch_dtype=torch.float32,
+    device_map="cpu",
 )
 
 
@@ -21,14 +19,12 @@ def respond(
     max_tokens,
     temperature,
     top_p,
-    hf_token: gr.OAuthToken,  # vẫn giữ cho đẹp, thực ra không xài
+    hf_token: gr.OAuthToken,  # giữ cho hợp với template, nhưng không dùng
 ):
-    # build messages theo format chat
     messages = [{"role": "system", "content": system_message}]
-    messages.extend(history)  # history đã là list[{"role","content"}]
+    messages.extend(history)
     messages.append({"role": "user", "content": message})
 
-    # dùng chat template nếu model có
     if hasattr(tokenizer, "apply_chat_template"):
         prompt = tokenizer.apply_chat_template(
             messages,
@@ -36,17 +32,9 @@ def respond(
             add_generation_prompt=True,
         )
     else:
-        # fallback tự ráp prompt đơn giản
-        parts = [f"System: {system_message}\n"]
-        for m in history:
-            if m["role"] == "user":
-                parts.append(f"User: {m['content']}\n")
-            elif m["role"] == "assistant":
-                parts.append(f"Assistant: {m['content']}\n")
-        parts.append(f"User: {message}\nAssistant:")
-        prompt = "".join(parts)
+        prompt = message
 
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(prompt, return_tensors="pt")
 
     outputs = model.generate(
         **inputs,
@@ -57,12 +45,11 @@ def respond(
         pad_token_id=tokenizer.eos_token_id,
     )
 
-    generated_ids = outputs[0, inputs.input_ids.shape[1]:]
-    full_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
+    gen_ids = outputs[0, inputs.input_ids.shape[1]:]
+    text = tokenizer.decode(gen_ids, skip_special_tokens=True)
 
-    # stream từng ký tự cho hợp với ChatInterface stream
     resp = ""
-    for ch in full_text:
+    for ch in text:
         resp += ch
         yield resp
 
@@ -86,9 +73,8 @@ chatbot = gr.ChatInterface(
 
 with gr.Blocks() as demo:
     with gr.Sidebar():
-        gr.LoginButton()  # nếu m muốn bắt user login HF mới xài
+        gr.LoginButton()
     chatbot.render()
 
-
 if __name__ == "__main__":
     demo.launch()