Update app.py

app.py

@@ -17,7 +17,7 @@ quantization_config = BitsAndBytesConfig(load_in_8bit=True)
 
 model = AutoModelForCausalLM.from_pretrained(
         MODEL,
-        torch_dtype=torch.bfloat16,
+        torch_dtype=torch.float16,
         device_map="auto",
         quantization_config=quantization_config)
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
@@ -61,7 +61,7 @@ def translate(
     print(f'Text is - {source_text}')
     
     prompt = Prompt_template(source_text, source_lang, target_lang)
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
         
     generate_kwargs = dict(
         input_ids=input_ids,
@@ -73,8 +73,10 @@ def translate(
     generate_ids = model.generate(**generate_kwargs)
     
     resp = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+
+    print(resp)
     
-    return resp
+    yield resp
 
 CSS = """
     h1 {