Update app.py

app.py

@@ -1,16 +1,19 @@
 import gradio as gr
-import transformers
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
 import re
 
 # Initialize the model
 model_id = "Detsutut/Igea-350M-v0.0.1"
-pipeline = transformers.pipeline(
+
+model = AutoModelForCausalLM.from_pretrained(model_id, load_in_8bit=True, device_map='auto')
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+
+gen_pipeline = pipeline(
     "text-generation",
-    model=model_id,
-    model_kwargs={"torch_dtype": torch.bfloat16},
-    load_in_8bit=True,
-    device_map="auto"
+    model=model,
+    tokenizer=tokenizer
 )
 
 # Define the function to generate text