Create app.py

app.py

@@ -0,0 +1,95 @@
+import os
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, logging as hf_logging
+from threading import Thread
+import gradio as gr
+from huggingface_hub import login
+
+# --- Hugging Face authentication ---
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if HF_TOKEN is None:
+    raise ValueError("Please set the HF_TOKEN environment variable.")
+login(token=HF_TOKEN)
+
+hf_logging.set_verbosity_error()  # suppress warnings
+
+# --- Model ID ---
+model_id = "motionlabs/NEWT-1.7B-QWEN-PREVIEW"
+
+# --- Logs helper ---
+log_messages = []
+
+def log(msg):
+    log_messages.append(msg)
+    print(msg)
+    return "\n".join(log_messages)
+
+log("Initializing tokenizer and model…")
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=HF_TOKEN)
+log("Tokenizer loaded.")
+
+# Load model
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    use_auth_token=HF_TOKEN
+)
+log("Model loaded.")
+
+# --- Chat streaming ---
+def stream_chat(history, message):
+    messages = []
+    for user, bot in history:
+        messages.append({"role": "user", "content": user})
+        if bot:
+            messages.append({"role": "assistant", "content": bot})
+    messages.append({"role": "user", "content": message})
+
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    gen_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=1024,
+        do_sample=True,
+        top_p=0.9,
+        temperature=0.7,
+    )
+
+    thread = Thread(target=model.generate, kwargs=gen_kwargs)
+    thread.start()
+
+    output_text = ""
+    for token in streamer:
+        output_text += token
+        yield history + [(message, output_text)]
+
+# --- Gradio UI ---
+with gr.Blocks(title=f"Chat with {model_id}") as demo:
+    gr.Markdown(f"# Chat with {model_id}")
+    
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(placeholder="Type your message here…")
+    clear = gr.Button("Clear")
+    logs = gr.Textbox(label="Logs", value="\n".join(log_messages), interactive=False)
+
+    def user_submit(user_message, history):
+        return "", history + [(user_message, None)]
+
+    msg.submit(user_submit, [msg, chatbot], [msg, chatbot]).then(
+        stream_chat, [chatbot, msg], chatbot
+    )
+    clear.click(lambda: None, None, chatbot, queue=False)
+
+demo.queue()
+demo.launch()