adjust model

app.py

@@ -1,27 +1,74 @@
 from openai import OpenAI
 import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
+from threading import Thread
 
-api_key = "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"
-client = OpenAI(api_key=api_key, base_url="https://api.perplexity.ai")
+#api_key = "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"
+#client = OpenAI(api_key=api_key, base_url="https://api.perplexity.ai")
 
-def predict(message, history):
-    history_openai_format = []
+#def predict(message, history):
+#    history_openai_format = []
     #history_openai_format.append({"role": "system", "content": "Answer the questions in japanese only"})
     
     #for human, assistant in history:
         #history_openai_format.append({"role": "user", "content": human })
         #history_openai_format.append({"role": "assistant", "content":assistant})
-    history_openai_format.append({"role": "user", "content": message})
+#    history_openai_format.append({"role": "user", "content": message})
   
-    response = client.chat.completions.create(model='sonar-small-chat',
-    messages= history_openai_format,
-    temperature=1.0,
-    stream=True)
+#    response = client.chat.completions.create(model='sonar-small-chat',
+#    messages= history_openai_format,
+#    temperature=1.0,
+#    stream=True)
+
+#    partial_message = ""
+#    for chunk in response:
+#        if chunk.choices[0].delta.content is not None:
+#              partial_message = partial_message + chunk.choices[0].delta.content
+#              yield partial_message
+
+#gr.ChatInterface(predict).launch()
+
+
+tokenizer = AutoTokenizer.from_pretrained("togethercomputer/RedPajama-INCITE-Chat-3B-v1")
+model = AutoModelForCausalLM.from_pretrained("togethercomputer/RedPajama-INCITE-Chat-3B-v1", torch_dtype=torch.float16)
+model = model.to('cuda:0')
+
+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("cuda")
+    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+    generate_kwargs = dict(
+        model_inputs,
+        streamer=streamer,
+        max_new_tokens=1024,
+        do_sample=True,
+        top_p=0.95,
+        top_k=1000,
+        temperature=1.0,
+        num_beams=1,
+        stopping_criteria=StoppingCriteriaList([stop])
+        )
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
 
     partial_message = ""
-    for chunk in response:
-        if chunk.choices[0].delta.content is not None:
-              partial_message = partial_message + chunk.choices[0].delta.content
-              yield partial_message
+    for new_token in streamer:
+        if new_token != '<':
+            partial_message += new_token
+            yield partial_message
 
 gr.ChatInterface(predict).launch()