Update app.py

app.py

@@ -9,17 +9,18 @@ question = "Qual é o maior planeta do sistema solar?"
 
 before = datetime.datetime.now()
 
-from transformers.modeling_outputs import Seq2SeqModelOutput, BaseModelOutput
+# from transformers.modeling_outputs import Seq2SeqModelOutput, BaseModelOutput
 
-from transformers import AutoTokenizer, XGLMModel
-import torch
+from mlx_lm import load, generate
+
+model, tokenizer = load("mlx-community/Meta-Llama-3.1-8B-Instruct-4bit")
 
 prompt = "Question: Qual é o maior planeta do sistema solar ?"
-tokenizer = AutoTokenizer.from_pretrained("facebook/xglm-564M",  use_fast=False)
-model = XGLMModel.from_pretrained("facebook/xglm-564M")
 
-inputs = tokenizer(prompt, return_tensors="pt")
-outputs = model(**inputs) #, labels=inputs["input_ids"])
+response = generate(model, tokenizer, prompt=prompt, verbose=True)
+
+# inputs = tokenizer(prompt, return_tensors="pt")
+# outputs = model(**inputs) #, labels=inputs["input_ids"])
 
 
 # last_hidden_states = outputs.last_hidden_state
@@ -32,8 +33,8 @@ outputs = model(**inputs) #, labels=inputs["input_ids"])
 
 # XGLMForCausalLM
 
-outputs = model(**inputs)
-output = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+# outputs = model(**inputs)
+# output = tokenizer.batch_decode(outputs, skip_special_tokens=True)
 
 
 # decoded = tokenizer.decode(output)
@@ -54,7 +55,7 @@ with st.container():
     st.write('\n\n')
     st.write('LLM-LANAChat\n\n')
     # st.write(outputs)
-    st.write(output)
+    st.write(response)
 
 print('\nsaida gerada.')
 print('\n\n')