Update app.py

app.py

@@ -9,52 +9,34 @@ question = "Qual é o maior planeta do sistema solar?"
 
 before = datetime.datetime.now()
 
-# from transformers.modeling_outputs import Seq2SeqModelOutput, BaseModelOutput
 
-from mlx_lm import load, generate
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model, tokenizer = load("mlx-community/Meta-Llama-3.1-8B-Instruct-4bit")
+model_id = "hugging-quants/Meta-Llama-3.1-8B-Instruct-BNB-NF4"
+prompt = [
+  {"role": "system", "content": "You are a helpful assistant"},
+  {"role": "user", "content": question},
+]
 
-prompt = "Question: Qual é o maior planeta do sistema solar ?"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-response = generate(model, tokenizer, prompt=prompt, verbose=True)
+inputs = tokenizer.apply_chat_template(prompt, tokenize=True, add_generation_prompt=True, return_tensors="pt").cuda()
 
-# inputs = tokenizer(prompt, return_tensors="pt")
-# outputs = model(**inputs) #, labels=inputs["input_ids"])
+model = AutoModelForCausalLM.from_pretrained(
+  model_id,
+  torch_dtype=torch.bfloat16,
+  low_cpu_mem_usage=True,
+  device_map="auto",
+)
 
+outputs = model.generate(inputs, do_sample=True, max_new_tokens=256)
+response = tokenizer.batch_decode(outputs, skip_special_tokens=True)
 
-# last_hidden_states = outputs.last_hidden_state
 
-# output = last_hidden_states #['last_hidden_states']
-
-
-# input_text = "The theory of special relativity states "
-# input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
-
-# XGLMForCausalLM
-
-# outputs = model(**inputs)
-# output = tokenizer.batch_decode(outputs, skip_special_tokens=True)
-
-
-# decoded = tokenizer.decode(output)
-
-# output = BaseModelOutput(last_hidden_states['last_hidden_states'])
-
-# logits = last_hidden_states.logits
-
-# output = last_hidden_states[0][0]
-
-# decoded = tokenizer.decode(output) # [0][0]
-# print(decoded)
-
-# output = Seq2SeqModelOutput(output)
-
-# output = tokenizer.batch_decode(output, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
 with st.container():
     st.write('\n\n')
     st.write('LLM-LANAChat\n\n')
-    # st.write(outputs)
     st.write(response)
 
 print('\nsaida gerada.')