Spaces:
Running
Running
File size: 1,988 Bytes
dea3a07 2fc7e1b a751c84 2fc7e1b dea3a07 2fc7e1b e604a26 af0df21 2fc7e1b af0df21 a751c84 2fc7e1b 32dbfef a751c84 2fc7e1b e604a26 a751c84 2fc7e1b a751c84 af0df21 2fc7e1b a751c84 2fc7e1b a751c84 cba2178 468f1f8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
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, share=True)
logger.info("end")
|