Update app.py

app.py

@@ -2,22 +2,17 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
-model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", device_map="auto")
+model_id = "mistralai/Mixtral-8x22B-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id)
 
 # Function to generate text using the model
 def generate_text(prompt, max_length=500, temperature=0.7, top_k=50, top_p=0.95, num_return_sequences=1):
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
-    output = model.generate(
-        input_ids,
-        max_length=max_length,
-        temperature=temperature,
-        top_k=top_k,
-        top_p=top_p,
-        num_return_sequences=num_return_sequences,
-    )
-    generated_text = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
-    return generated_text
+    text = prompt
+    inputs = tokenizer(text, return_tensors="pt")
+    
+    outputs = model.generate(**inputs, max_new_tokens=20)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # Create the Gradio interface
 iface = gr.Interface(