app.py

import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import BitsAndBytesConfig

# Function to load a quantized model
def load_quantized_model():
    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
    config = BitsAndBytesConfig.from_dict({"load_in_4bit": True})
    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B-Instruct", quantization_config=config)
    return model, tokenizer

model, tokenizer = load_quantized_model()

# Simple prediction function for Gradio
def generate_response(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(**inputs)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Gradio interface
iface = gr.Interface(
    fn=generate_response,
    inputs="text",
    outputs="text",
    title="Quantized Model Chatbot"
)

iface.launch()