Update app.py

app.py

@@ -1,18 +1,54 @@
-# Load model directly
+import gradio as gr
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-tokenizer = AutoTokenizer.from_pretrained("LivSterling/rc-tutor-llama3-merged")
-model = AutoModelForCausalLM.from_pretrained("LivSterling/rc-tutor-llama3-merged")
-messages = [
-    {"role": "user", "content": "Who are you?"},
-]
-inputs = tokenizer.apply_chat_template(
-	messages,
-	add_generation_prompt=True,
-	tokenize=True,
-	return_dict=True,
-	return_tensors="pt",
-).to(model.device)
-
-outputs = model.generate(**inputs, max_new_tokens=40)
-print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:]))
+MODEL_ID = "LivSterling/rc-tutor-llama3-merged"
+
+# Load model & tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float16,
+    device_map="auto",
+)
+
+def chat_fn(message, history):
+    # Convert Gradio history to chat template format
+    messages = []
+    for user_msg, bot_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "assistant", "content": bot_msg})
+
+    messages.append({"role": "user", "content": message})
+
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_tensors="pt",
+    ).to(model.device)
+
+    outputs = model.generate(
+        inputs,
+        max_new_tokens=256,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9,
+    )
+
+    response = tokenizer.decode(
+        outputs[0][inputs.shape[-1]:],
+        skip_special_tokens=True,
+    )
+
+    return response
+
+# Gradio UI
+demo = gr.ChatInterface(
+    fn=chat_fn,
+    title="RC Tutor LLaMA-3 Chatbot",
+    description="Powered by LivSterling/rc-tutor-llama3-merged",
+)
+
+if __name__ == "__main__":
+    demo.launch()