Update app.py with working Gradio chatbot interface

app.py

@@ -1,13 +1,17 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+from transformers import AutoTokenizer
+from transformers import AutoModelForCausalLM
 import gradio as gr
 
-# Load tokenizer and model with trust_remote_code=True
-tokenizer = AutoTokenizer.from_pretrained("BWComedian/CSM-1B", trust_remote_code=True)
+# Load LLaMA tokenizer (base tokenizer)
+tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
+
+# Load your CSM-1B model (must support causal LM interface)
 model = AutoModelForCausalLM.from_pretrained("BWComedian/CSM-1B", trust_remote_code=True)
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
+model.eval()
 
 def generate_response(prompt):
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
@@ -31,4 +35,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()