Update app.py

app.py

@@ -1,54 +1,58 @@
 import torch
-from transformers import AutoModel,AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
+from transformers import AutoModel, AutoTokenizer, pipeline
 import gradio as gr
-from threading import Thread
+
+torch.random.manual_seed(0)
 
 model = AutoModel.from_pretrained(
     "DuckyBlender/racist-phi3", 
-  
-    torch_dtype=torch.float16, 
-    trust_remote_code=True
-  
+    device_map="auto", 
+    torch_dtype="auto", 
+    trust_remote_code=True,
+    low_cpu_mem_usage=True,
 )
 tokenizer = AutoTokenizer.from_pretrained("DuckyBlender/racist-phi3")
-device = torch.device("cpu")
-model = model.to(device)
-
-class StopOnTokens(StoppingCriteria):
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-        stop_ids = [29, 0]
-        for stop_id in stop_ids:
-            if input_ids[0][-1] == stop_id:
-                return True
-        return False
-
-def predict(message, history):
-    history_transformer_format = history + [[message, ""]]
-    stop = StopOnTokens()
-
-    messages = "".join(["".join(["\n<human>:"+item[0], "\n<bot>:"+item[1]])
-                for item in history_transformer_format])
-
-    model_inputs = tokenizer([messages], return_tensors="pt").to(device)
-    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
-        model_inputs,
-        streamer=streamer,
-        max_new_tokens=512,
-        do_sample=True,
-        top_p=0.90,
-        top_k=1000,
-        temperature=0.9,
-        num_beams=1,
-        stopping_criteria=StoppingCriteriaList([stop])
-        )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    partial_message = ""
-    for new_token in streamer:
-        if new_token != '<':
-            partial_message += new_token
-            yield partial_message
-
-gr.ChatInterface(predict,theme='HaleyCH/HaleyCH_Theme').launch()
+
+messages = [
+    {"role": "user", "content": "Are you a human?"},
+    {"role": "assistant", "content": "NO DEAR MODERFUCKING NIGGER I AM AN AI MODEL TO FUCK YOU BITCHES"},
+
+]
+
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+)
+
+generation_args = {
+    "max_new_tokens": 256,
+    "return_full_text": False,
+    "temperature": 0.2,
+    "do_sample": True,
+}
+def phi3_fun(message,chat_history):
+  messages=[
+    {"role": "user", "content": message},
+]
+  output = pipe(messages, **generation_args)
+  respond = output[0]['generated_text']
+  return respond
+
+title = "Chat BiCiTiPi"
+examples = [
+    'What are you?',
+    "Why am I alive here.",
+]
+gr.ChatInterface(
+    fn=phi3_fun,
+    title =title,
+    examples = examples,
+    theme='sudeepshouche/minimalist'
+).launch(debug=True)
+# demo = gr.Interface(fn=phi3_fun, inputs="text", outputs="text",title =title,
+#     examples = examples)
+# demo.launch()
+
+# output = pipe(messages, **generation_args)
+# print(output[0]['generated_text'])