Update app.py

app.py

@@ -8,8 +8,10 @@ from logging.handlers import RotatingFileHandler
 import torch
 import spaces
 import gradio as gr
-from transformers import AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
 from langchain_huggingface import ChatHuggingFace
+from langchain.prompts import PromptTemplate
+from langchain.chains import LLMChain
 
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
@@ -33,34 +35,42 @@ COLOR = "blue"
 EMOJI = "🤖"
 DESCRIPTION = f"This is the {MODEL_NAME} model designed for coding assistance and general AI tasks."
 
+# Prompt template for conversation
+template = """<|im_start|>system
+{system_prompt}
+<|im_end|>
+{history}
+<|im_start|>user
+{human_input}
+<|im_end|>
+<|im_start|>assistant
+"""
+prompt = PromptTemplate(template=template, input_variables=["system_prompt", "history", "human_input"])
+
+# Format the conversation history
+def format_history(history):
+    formatted = ""
+    for human, ai in history:
+        formatted += f"<|im_start|>user\n{human}\n<|im_end|>\n<|im_start|>assistant\n{ai}\n<|im_end|>\n"
+    return formatted
+
 @spaces.GPU()
 def predict(message, history, system_prompt, temperature, max_new_tokens, top_k, repetition_penalty, top_p):
     logger.debug(f"Received prediction request: message='{message}', system_prompt='{system_prompt}'")
-    if CHAT_TEMPLATE == "Auto":
-        instruction = system_prompt + "\n\n"
-        for user, assistant in history:
-            instruction += f"User: {user}\nAssistant: {assistant}\n"
-        instruction += f"User: {message}\nAssistant:"
-    elif CHAT_TEMPLATE == "ChatML":
-        instruction = '<|im_start|>system\n' + system_prompt + '\n<|im_end|>\n'
-        for user, assistant in history:
-            instruction += f'<|im_start|>user\n{user}\n<|im_end|>\n<|im_start|>assistant\n{assistant}\n<|im_end|>\n'
-        instruction += f'<|im_start|>user\n{message}\n<|im_end|>\n<|im_start|>assistant\n'
-    elif CHAT_TEMPLATE == "Mistral Instruct":
-        instruction = f'<s>[INST] {system_prompt}\n'
-        for user, assistant in history:
-            instruction += f'{user} [/INST] {assistant}</s>[INST]'
-        instruction += f' {message} [/INST]'
-    else:
-        raise Exception("Incorrect chat template, select 'Auto', 'ChatML' or 'Mistral Instruct'")
-    print(instruction)
     
-    response = chat_model.predict(instruction)
-    outputs = []
+    formatted_history = format_history(history)
+    
+    chat_model.temperature = temperature
+    chat_model.max_new_tokens = max_new_tokens
+    chat_model.top_k = top_k
+    chat_model.repetition_penalty = repetition_penalty
+    chat_model.top_p = top_p
+
+    chain = LLMChain(llm=chat_model, prompt=prompt)
+    
     try:
-        for token in response:
-            outputs.append(token)
-            yield "".join(outputs)
+        for chunk in chain.stream({"system_prompt": system_prompt, "history": formatted_history, "human_input": message}):
+            yield chunk["text"]
         logger.debug(f"Prediction completed successfully for message: '{message}'")
     except Exception as e:
         logger.exception(f"Error during prediction for message '{message}': {str(e)}")
@@ -75,12 +85,12 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
 chat_model = ChatHuggingFace(
     model_name=MODEL_ID,
-    huggingfacehub_api_token=os.environ["HF_TOKEN"],
     model_kwargs={
         "device_map": "auto",
         "quantization_config": quantization_config,
         "attn_implementation": "flash_attention_2",
-    }
+    },
+    tokenizer=tokenizer
 )
 
 logger.debug("Model and tokenizer loaded successfully")
@@ -101,7 +111,6 @@ gr.ChatInterface(
        ["How do I find the maximum element in an array using Kotlin?"],
        ["Write a Rust program to generate the Fibonacci sequence up to the 10th number."]
     ],
-    additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False),
     additional_inputs=[
         gr.Textbox("You are a code assistant.", label="System prompt"),
         gr.Slider(0, 1, 0.3, label="Temperature"),