app.py

import gradio as gr
from unsloth import FastLanguageModel
from transformers import AutoTokenizer, TextStreamer

# Load the model and tokenizer
model_name = "Rafay17/Llama3.2_1b_customModel2"  # Your custom model
model, tokenizer = FastLanguageModel.from_pretrained(model_name)
FastLanguageModel.for_inference(model)  # Enable the model for inference

# Function to generate a response
def generate_response(message, history, max_tokens, temperature, top_p):
    # Prepare the labeled prompt for response generation
    labeled_prompt = f"User Input: {message}\nResponse:"

    # Tokenize the input
    inputs = tokenizer(
        [labeled_prompt],
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=512,
    ).to("cuda")

    # Generate the response
    text_streamer = TextStreamer(tokenizer, skip_prompt=True)
    response = ""
    for token in model.generate(
        input_ids=inputs.input_ids,
        attention_mask=inputs.attention_mask,
        streamer=text_streamer,
        max_new_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        pad_token_id=tokenizer.eos_token_id,
    ):
        response += token

    return response


# Define the Gradio interface
demo = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter your message here..."),
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=512, value=64, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.9, label="Top-p (nucleus sampling)"),
    ],
    outputs=gr.Textbox(label="Chatbot Response"),
    live=True
)

if __name__ == "__main__":
    demo.launch()