Update app.py

app.py

@@ -1,10 +1,26 @@
-from transformers import pipeline as pl
 import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-oracle = pl("stabilityai/stablelm-2-zephyr-1_6b")
+tokenizer = AutoTokenizer.from_pretrained('stabilityai/stablelm-2-zephyr-1_6b')
+model = AutoModelForCausalLM.from_pretrained(
+    'stabilityai/stablelm-2-zephyr-1_6b',
+    device_map="auto"
+)
 
 def pipe(text: str):
-    return oracle(text)
+    tokens = model.generate(
+        inputs.to(model.device),
+        max_new_tokens=1024,
+        temperature=0.5,
+        do_sample=True
+    )
+
+    inputs = tokenizer.apply_chat_template(
+        text,
+        add_generation_prompt=True,
+        return_tensors='pt'
+    )
+    return tokenizer.decode(tokens[0], skip_special_tokens=False)
 
 if __name__ == "__main__":
     interface = gr.Interface(pipe, gr.Textbox(label="Prompt"), gr.Textbox(label="Response"), title="Text Completion")