app.py

import os
import json
import subprocess
from threading import Thread
import logging
from logging.handlers import RotatingFileHandler

import torch
import spaces
import gradio as gr
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)

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)
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
"""
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}'")
    
    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 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)}")
        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,
    model_kwargs={
        "device_map": "auto",
        "quantization_config": quantization_config,
        "attn_implementation": "flash_attention_2",
    },
    tokenizer=tokenizer
)

logger.debug("Model and tokenizer loaded successfully")

gr.ChatInterface(
    predict,
    title=EMOJI + " " + MODEL_NAME,
    description=DESCRIPTION,
    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."],
       ["How do I implement a binary search algorithm in C?"],
       ["Write a Ruby script to read a file and count the number of lines in it."],
       ["Create a Swift class to represent a bank account with deposit and withdrawal methods."],
       ["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=[
        gr.Textbox("You are a code assistant.", label="System prompt"),
        gr.Slider(0, 1, 0.3, label="Temperature"),
        gr.Slider(128, 4096, 1024, label="Max new tokens"),
        gr.Slider(1, 80, 40, label="Top K sampling"),
        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),
).queue().launch()

logger.debug("Chat interface initialized and launched")