Update app.py

app.py

@@ -1,7 +1,44 @@
 import gradio as gr
 
+from transformers import AutoModelForCausalLM, AutoTokenizer, StopStringCriteria, StoppingCriteriaList
+import torch
+
+# Load the tokenizer and model
+repo_name = "nvidia/Hymba-1.5B-Instruct"
+
+tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(repo_name, trust_remote_code=True)
+model = model.cuda().to(torch.bfloat16)
+
+# Chat with Hymba
+# prompt = input()
+prompt = "Who are you?"
+
+messages = [
+    {"role": "system", "content": "You are a helpful assistant."}
+]
+messages.append({"role": "user", "content": prompt})
+
+# Apply chat template
+tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to('cuda')
+stopping_criteria = StoppingCriteriaList([StopStringCriteria(tokenizer=tokenizer, stop_strings="</s>")])
+outputs = model.generate(
+    tokenized_chat, 
+    max_new_tokens=256,
+    do_sample=False,
+    temperature=0.7,
+    use_cache=True,
+    stopping_criteria=stopping_criteria
+)
+input_length = tokenized_chat.shape[1]
+response = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
+
+
+
 def greet(name):
-    return "Hello " + name + "!!"
+    print(f"User: prompt")
+    print(f"Model response: {response}")
+    # return "Hello " + name + "!!"
 
 demo = gr.Interface(fn=greet, inputs="text", outputs="text")
 demo.launch()