Update app.py

app.py

@@ -2,7 +2,7 @@ import torch
 from PIL import Image
 import gradio as gr
 import spaces
-from transformers import LlamaForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
 import os
 from threading import Thread
 from polyglot.detect import Detector
@@ -15,7 +15,7 @@ TITLE = "<h1><center>LLaMAX3-8B-Translation</center></h1>"
 
 quantization_config = BitsAndBytesConfig(load_in_8bit=True)
 
-model = LlamaForCausalLM.from_pretrained(
+model = AutoModelForCausalLM.from_pretrained(
         MODEL,
         torch_dtype=torch.bfloat16,
         device_map="auto",
@@ -61,12 +61,12 @@ def translate(
     print(f'Text is - {source_text}')
     
     prompt = Prompt_template(source_text, source_lang, target_lang)
-    inputs = tokenizer(prompt, return_tensors="pt")
+    input_ids = tokenizer(prompt, return_tensors="pt").to(model.device)
         
     streamer = TextIteratorStreamer(tokenizer, **{"skip_special_tokens": True, "skip_prompt": True, 'clean_up_tokenization_spaces':False,})
 
     generate_kwargs = dict(
-        inputs.input_ids, 
+        input_ids=input_ids,
         streamer=streamer,
         max_length=max_length, 
         do_sample=True,