Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 import random
-from transformers import AutoTokenizer,AutoModelForCausalLM
+from transformers import AutoTokenizer
 from mySystemPrompt import SYSTEM_PROMPT
 
 
@@ -10,7 +10,6 @@ checkpoint = "CohereForAI/c4ai-command-r-plus"
 # Inference client with the model (And HF-token if needed)
 client = InferenceClient(checkpoint)
 tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-model = AutoModelForCausalLM.from_pretrained(checkpoint)
 # Tokenizer chat template correction(Only works for mistral models)
 #chat_template = open("mistral-instruct.jinja").read()
 #chat_template = chat_template.replace('    ', '').replace('\n', '')
@@ -22,9 +21,9 @@ def format_prompt(message,chatbot,system_prompt):
         messages.append({"role": "user", "content":user_message})
         messages.append({"role": "assistant", "content":bot_message})
     messages.append({"role": "user", "content":message})
-    #newPrompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, return_tensors="pt")
-    #print(newPrompt)
-    newPrompt = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
+    newPrompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, return_tensors="pt")
+    print(newPrompt)
+    #newPrompt = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
     return newPrompt
 
 def inference(message, history, systemPrompt=SYSTEM_PROMPT, temperature=0.9, maxTokens=512, topP=0.9, repPenalty=1.1):
@@ -41,18 +40,14 @@ def inference(message, history, systemPrompt=SYSTEM_PROMPT, temperature=0.9, max
         seed=random.randint(0, 999999999),
     )
     # Generating the response by passing the prompt in right format plus the client settings
-    #stream = client.text_generation(format_prompt(message, history, systemPrompt),
-      #                      **client_settings)
+    stream = client.text_generation(format_prompt(message, history, systemPrompt),
+                            **client_settings)
     # Reading the stream
-    #partial_response = ""
-    #for stream_part in stream:
-     #   partial_response += stream_part.token.text
-      #  yield partial_response
-    gen_tokens = model.generate(
-    format_prompt(message,history,systemPrompt),
-    **client_settings)
-    output = tokenizer.decode(gen_tokens[0])
-    return output
+    partial_response = ""
+    for stream_part in stream:
+        partial_response += stream_part.token.text
+        yield partial_response
+