app.py

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

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

model_name = "google/gemma-2-2b-it"
try:
    logger.info(f"Loading model: {model_name}")
    tokenizer = AutoTokenizer.from_pretrained(model_name, token=os.getenv("HF_TOKEN"))
    use_gpu = torch.cuda.is_available()
    logger.info(f"GPU available: {use_gpu}")
    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True
    ) if use_gpu else None
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        token=os.getenv("HF_TOKEN"),
        low_cpu_mem_usage=True,
        quantization_config=quantization_config
    )
    logger.info("Model loaded successfully")
except Exception as e:
    logger.error(f"Model load error: {e}")
    raise

def generate_text(text, max_length=50):
    try:
        logger.info(f"Generating text for input: {text}")
        inputs = tokenizer(text, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
        outputs = model.generate(**inputs, max_length=max_length)
        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
        logger.info(f"Generated text: {result}")
        return result
    except Exception as e:
        logger.error(f"Generation error: {e}")
        return f"Error: {str(e)}"

iface = gr.Interface(
    fn=generate_text,
    inputs=[gr.Textbox(label="Input Text"), gr.Slider(10, 100, value=50, label="Max Length")],
    outputs=gr.Textbox(label="Generated Text"),
    title="Gemma 2 API"
)

if __name__ == "__main__":
    logger.info("Launching Gradio interface")
    iface.launch(server_name="0.0.0.0", server_port=8080)