Spaces:
Paused
Paused
File size: 1,827 Bytes
8f86f02 4a9540b 258e5e7 4a9540b |
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 |
import torch
from transformers import AutoProcessor, AutoModelForImageTextToText, TextStreamer
from peft import PeftModel
import gradio as gr
# Load base model and processor
base_model_id = "unsloth/gemma-3-12b-it-unsloth-bnb-4bit"
adapter_model_id = "adarsh3601/my_gemma_pt3"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
processor = AutoProcessor.from_pretrained(base_model_id)
model = AutoModelForImageTextToText.from_pretrained(base_model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, device_map="auto")
# Apply adapter (LoRA)
model = PeftModel.from_pretrained(model, adapter_model_id)
model.eval()
streamer = TextStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
# Helper to format messages using the chat template
def format_chat(messages):
formatted = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
return formatted
# Chat function
def chat(message, history):
messages = []
# Format history into messages
for user_msg, bot_msg in history:
messages.append({"role": "user", "content": user_msg})
messages.append({"role": "assistant", "content": bot_msg})
messages.append({"role": "user", "content": message})
prompt = format_chat(messages)
inputs = processor(prompt, return_tensors="pt").to(device)
with torch.no_grad():
outputs = model.generate(**inputs, max_new_tokens=512, streamer=streamer)
decoded = processor.batch_decode(outputs, skip_special_tokens=True)[0]
response = decoded.split("<end_of_turn>")[0].strip().split("<start_of_turn>model")[-1].strip()
return response
# Gradio interface
gui = gr.ChatInterface(fn=chat, title="Gemma-3 Chatbot", description="Fine-tuned on adarsh3601/my_gemma_pt3")
gui.launch()
|