Update app.py

app.py

@@ -1,16 +1,16 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelWithLMHead
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
 @st.cache_data
 def load_model(model_name):
-    model = AutoModelWithLMHead.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
+    model = AutoModelForSeq2SeqLM.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
     return model
 
 model = load_model("mistralai/Mixtral-8x7B-Instruct-v0.1")
 
 def infer(input_ids, max_length, temperature, top_k, top_p):
-
+    input_ids = torch.tensor(input_ids).unsqueeze(0)
     output_sequences = model.generate(
         input_ids=input_ids,
         max_length=max_length,
@@ -20,8 +20,8 @@ def infer(input_ids, max_length, temperature, top_k, top_p):
         do_sample=True,
         num_return_sequences=1
     )
-
     return output_sequences
+
 default_value = "Ask me anything!"
 
 #prompts
@@ -39,28 +39,9 @@ if encoded_prompt.size()[-1] == 0:
 else:
     input_ids = encoded_prompt
 
-
 output_sequences = infer(input_ids, max_length, temperature, top_k, top_p)
 
-
-
 for generated_sequence_idx, generated_sequence in enumerate(output_sequences):
-    print(f"=== GENERATED SEQUENCE {generated_sequence_idx + 1} ===")
     generated_sequences = generated_sequence.tolist()
-
-    # Decode text
     text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)
-
-    # Remove all text after the stop token
-    #text = text[: text.find(args.stop_token) if args.stop_token else None]
-
-    # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
-    total_sequence = (
-        sent + text[len(tokenizer.decode(encoded_prompt[0], clean_up_tokenization_spaces=True)) :]
-    )
-
-    generated_sequences.append(total_sequence)
-    print(total_sequence)
-
-
-st.write(generated_sequences[-1])
+    st.write(text)