Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr 
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 sys_message = """
 This model can generate untruths, lies or inappropriate things. Only for testing and validation. 
@@ -13,22 +13,28 @@ generation_config = {
     "top_p": 0.95,
     "top_k": 30,
     "repetition_penalty": 1.1,
-    "stop_token": '### Instruction:'
+    "eos_token_id": pipe.tokenizer.eos_token_id 
 }
 
-tokenizer = AutoTokenizer.from_pretrained("tevykuch/sftsl0th")
-llm = AutoModelForCausalLM.from_pretrained("tevykuch/sftsl0th")
 
-def stream(prompt):
-    # Tokenize the prompt
-    inputs = tokenizer.encode(prompt, return_tensors="pt")
-    # Generate a response
-    output_ids = llm.generate(inputs, **generation_config)
-    # Decode the generated ids to a string
-    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return response
+pipe = pipeline("text-generation", model="tevykuch/sftsl0th", device=0, framework="pt")
+
+# tokenizer = AutoTokenizer.from_pretrained("tevykuch/sftsl0th")
+# llm = AutoModelForCausalLM.from_pretrained("tevykuch/sftsl0th")
 
+# def stream(prompt):
+#     # Tokenize the prompt
+#     inputs = tokenizer.encode(prompt, return_tensors="pt")
+#     # Generate a response
+#     output_ids = llm.generate(inputs, **generation_config)
+#     # Decode the generated ids to a string
+#     response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+#     return response
 
+def stream(prompt):
+    outputs = pipe(prompt, **generation_config)
+    response = outputs[0]["generated_text"]
+    return response
 
 chat_interface = gr.ChatInterface(
     fn=stream,