BICORP commited on
Commit
f0be7cb
Β·
verified Β·
1 Parent(s): 1513f4e

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +31 -0
app.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from transformers import AutoTokenizer, AutoModelForCausalLM
4
+
5
+ # Model and Tokenizer Setup
6
+ model_name = "unsloth/gemma-3-4b-it-unsloth-bnb-4bit"
7
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
8
+ model = AutoModelForCausalLM.from_pretrained(
9
+ model_name,
10
+ load_in_4bit=True,
11
+ device_map="auto",
12
+ torch_dtype=torch.bfloat16, #important for speed.
13
+ )
14
+
15
+ def generate_response(prompt):
16
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
17
+ outputs = model.generate(**inputs, max_new_tokens=256) # Adjust max_new_tokens as needed
18
+ response = tokenizer.decode(outputs[0], skip_special_tokens=True)
19
+ return response
20
+
21
+ # Gradio Interface
22
+ iface = gr.Interface(
23
+ fn=generate_response,
24
+ inputs=gr.Textbox(lines=5, placeholder="Enter your prompt here..."),
25
+ outputs=gr.Textbox(),
26
+ title="Gemma 3-4B Inference",
27
+ description="Run the unsloth/gemma-3-4b-it-unsloth-bnb-4bit model.",
28
+ )
29
+
30
+ if __name__ == "__main__":
31
+ iface.launch(server_name="0.0.0.0", server_port=7860) #important for spaces.