Update app.py

app.py

@@ -4,7 +4,7 @@ import gradio as gr
 import os
 import spaces
 import torch
-from transformers import AutoModel, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from threading import Thread
 
 DESCRIPTION = '''
@@ -45,7 +45,7 @@ h1 {
 
 # Load the tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("alfredplpl/sarashina2-7b-it")
-model = AutoModel.from_pretrained("alfredplpl/sarashina2-7b-it", torch_dtype=torch.bfloat16)
+model = AutoModelForCausalLM.from_pretrained("alfredplpl/sarashina2-7b-it", torch_dtype=torch.bfloat16)
 model=model.to("cuda:0")
 
 @spaces.GPU()