Update app.py

app.py

@@ -1,9 +1,12 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 import spaces
 
-# Initialize the text generation pipeline outside the function for efficiency
-pipe = pipeline("text-generation", model="deepapaikar/LlamaKatz-3x8B")
+# Load model and tokenizer only once, outside the function
+model_name = "deepapaikar/LlamaKatz-3x8B"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name,  device_map='auto')
 
 @spaces
 def generate_text(input_text):
@@ -15,11 +18,10 @@ def generate_text(input_text):
     Returns:
         str: The generated text.
     """
-    messages = [
-        {"role": "user", "content": input_text},
-    ]
-    output = pipe(messages)
-    return output[0]['generated_text']  # Extract the generated text 
+    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs)
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return generated_text
 
 iface = gr.Interface(
     fn=generate_text,