Update app.py

app.py

@@ -1,18 +1,21 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+model_name = "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF"
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
+prompt = "How many r in strawberry?"
+messages = [{"role": "user", "content": prompt}]
 
-client = InferenceClient("nvidia/Llama-3.1-Nemotron-70B-Instruct-HF")
+tokenized_message = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True)
+response_token_ids = model.generate(tokenized_message['input_ids'].cuda(),attention_mask=tokenized_message['attention_mask'].cuda(),  max_new_tokens=4096, pad_token_id = tokenizer.eos_token_id)
+generated_tokens =response_token_ids[:, len(tokenized_message['input_ids'][0]):]
+generated_text = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+print(generated_text)
 
+# See response at top of model card
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
     messages = [{"role": "system", "content": system_message}]
 
     for val in history: