Update app.py

app.py

@@ -2,7 +2,6 @@ import os
 import logging
 from logging.handlers import RotatingFileHandler
 
-import torch
 import gradio as gr
 from transformers import AutoTokenizer, BitsAndBytesConfig
 from langchain_huggingface import ChatHuggingFace
@@ -14,45 +13,23 @@ log_file = '/tmp/app_debug.log'
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 file_handler = RotatingFileHandler(log_file, maxBytes=10*1024*1024, backupCount=5)
-file_handler.setLevel(logging.DEBUG)
-formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-file_handler.setFormatter(formatter)
+file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
 logger.addHandler(file_handler)
 
 logger.debug("Application started")
 
 MODEL_ID = "Qwen/Qwen2.5-Coder-7B-Instruct"
-CHAT_TEMPLATE = "ChatML"
 MODEL_NAME = MODEL_ID.split("/")[-1]
-CONTEXT_LENGTH = 16000
 
-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
-"""
+template = """<|im_start|>system\n{system_prompt}\n<|im_end|>\n{history}<|im_start|>user\n{human_input}\n<|im_end|>\n<|im_start|>assistant\n"""
 prompt = PromptTemplate(template=template, input_variables=["system_prompt", "history", "human_input"])
 
 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
+    return "".join([f"<|im_start|>user\n{h[0]}\n<|im_end|>\n<|im_start|>assistant\n{h[1]}\n<|im_end|>\n" for h in history])
 
 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}'")
     
-    formatted_history = format_history(history)
-    
     chat_model.temperature = temperature
     chat_model.max_new_tokens = max_new_tokens
     chat_model.top_k = top_k
@@ -62,41 +39,36 @@ def predict(message, history, system_prompt, temperature, max_new_tokens, top_k,
     chain = LLMChain(llm=chat_model, prompt=prompt)
     
     try:
+        formatted_history = format_history(history)
         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)}")
+        logger.exception(f"Error during prediction: {str(e)}")
         yield "An error occurred during processing."
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-quantization_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.bfloat16
-)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-
 chat_model = ChatHuggingFace(
     model_name=MODEL_ID,
+    tokenizer=tokenizer,
     model_kwargs={
         "device_map": "auto",
-        "quantization_config": quantization_config,
-    },
-    tokenizer=tokenizer
+        "quantization_config": BitsAndBytesConfig(load_in_4bit=True),
+    }
 )
 
 logger.debug("Model and tokenizer loaded successfully")
 
 gr.ChatInterface(
     predict,
-    title=EMOJI + " " + MODEL_NAME,
-    description=DESCRIPTION,
+    title=f"🤖 {MODEL_NAME}",
+    description=f"This is the {MODEL_NAME} model designed for coding assistance and general AI tasks.",
     examples=[
-       ["Can you solve the equation 2x + 3 = 11 for x in Python?"],
-       ["Write a Java program that checks if a number is even or odd."],
-       ["How can I reverse a string in JavaScript?"],
-       ["Create a C++ function to find the factorial of a number."],
-       ["Write a Python list comprehension to generate a list of squares of numbers from 1 to 10."],
+        ["Can you solve the equation 2x + 3 = 11 for x in Python?"],
+        ["Write a Java program that checks if a number is even or odd."],
+        ["How can I reverse a string in JavaScript?"],
+        ["Create a C++ function to find the factorial of a number."],
+        ["Write a Python list comprehension to generate a list of squares of numbers from 1 to 10."],
     ],
     additional_inputs=[
         gr.Textbox("You are a code assistant.", label="System prompt"),
@@ -106,7 +78,7 @@ gr.ChatInterface(
         gr.Slider(0, 2, 1.1, label="Repetition penalty"),
         gr.Slider(0, 1, 0.95, label="Top P sampling"),
     ],
-    theme=gr.themes.Soft(primary_hue=COLOR),
+    theme=gr.themes.Soft(primary_hue="blue"),
 ).queue().launch()
 
 logger.debug("Chat interface initialized and launched")