Update app.py

app.py

@@ -1,91 +1,21 @@
+
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
 import gradio as gr
-from gradio import ChatMessage
-from typing import Iterator
 
 checkpoint = "EpistemeAI/metatune-gpt20b-R0"
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load model + tokenizer
+device = "cuda"  # "cuda" or "cpu"
 tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-model = AutoModelForCausalLM.from_pretrained(
-    checkpoint,
-    torch_dtype="auto",
-).to(device)
-
-def format_history_for_model(messages):
-    """Convert the message list into a single string prompt"""
-    chat_prompt = ""
-    for msg in messages:
-        role = msg["role"]
-        content = msg["content"]
-        if role == "user":
-            chat_prompt += f"User: {content}\n"
-        else:
-            chat_prompt += f"Assistant: {content}\n"
-    return chat_prompt.strip()
-
-def stream_response(user_message: str, messages: list) -> Iterator[list]:
-    try:
-        print(f"User: {user_message}")
-        prompt = format_history_for_model(messages) + f"\nUser: {user_message}\nAssistant:"
-
-        # Tokenize
-        inputs = tokenizer(prompt, return_tensors="pt").to(device)
-        
-        # Stream output tokens
-        generated = model.generate(
-            **inputs,
-            max_new_tokens=256,
-            temperature=0.7,
-            do_sample=True,
-            top_p=0.9,
-            repetition_penalty=1.1,
-            pad_token_id=tokenizer.eos_token_id,
-        )
-        output_text = tokenizer.decode(generated[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
-        
-        # Send back message
-        messages.append(ChatMessage(role="assistant", content=output_text))
-        yield messages
-
-    except Exception as e:
-        messages.append(ChatMessage(role="assistant", content=f"Error: {str(e)}"))
-        yield messages
-
-def user_message(msg: str, history: list):
-    history.append(ChatMessage(role="user", content=msg))
-    return "", history
-
-# --- UI ---
-with gr.Blocks(theme=gr.themes.Citrus(), fill_height=True) as demo:
-    gr.Markdown("# Chat with Metatune GPT 20B 💭")
-
-    chatbot = gr.Chatbot(type="messages", label="Metatune 20B Chatbot", render_markdown=True)
-    with gr.Row():
-        input_box = gr.Textbox(label="Message", placeholder="Type your message here...")
-        clear_button = gr.Button("Clear")
-
-    msg_store = gr.State("")
+model = AutoModelForCausalLM.from_pretrained(checkpoint,torch_dtype="auto").to(device)
 
-    input_box.submit(
-        lambda msg: (msg, msg, ""),
-        inputs=[input_box],
-        outputs=[msg_store, input_box, input_box],
-        queue=False,
-    ).then(
-        user_message,
-        inputs=[msg_store, chatbot],
-        outputs=[input_box, chatbot],
-        queue=False,
-    ).then(
-        stream_response,
-        inputs=[msg_store, chatbot],
-        outputs=chatbot,
-    )
+def predict(message, history):
+    history.append({"role": "user", "assistant": message})
+    input_text = tokenizer.apply_chat_template(history, tokenize=False)
+    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)  
+    outputs = model.generate(inputs, max_new_tokens=3200, temperature=0.2, top_p=0.9, do_sample=True)
+    decoded = tokenizer.decode(outputs[0])
+    response = decoded.split("<|im_start|>assistant\n")[-1].split("<|im_end|>")[0]
+    return response
 
-    clear_button.click(lambda: ([], "", ""), outputs=[chatbot, input_box, msg_store])
+demo = gr.ChatInterface(predict, type="messages")
 
-if __name__ == "__main__":
-    demo.launch(debug=True)
+demo.launch()