kimhyunwoo commited on
Commit
8f6729a
·
verified ·
1 Parent(s): c5ec987

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +60 -114
app.py CHANGED
@@ -1,118 +1,64 @@
1
- from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, logging
2
- from huggingface_hub import login
3
- import torch
4
- import os
5
  import gradio as gr
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6
 
7
- # --- 1. Authentication (Using User-Provided Token) ---
8
-
9
- def authenticate(token):
10
- """Attempts to authenticate with the provided token."""
11
- try:
12
- login(token=token)
13
- return True
14
- except Exception as e:
15
- print(f"Authentication failed: {e}")
16
- return False
17
-
18
- # --- 2. Model and Tokenizer Setup ---
19
-
20
- def load_model_and_tokenizer(model_name="google/gemma-3-1b-it"):
21
- """Loads the model and tokenizer."""
22
- try:
23
- logging.set_verbosity_error()
24
- tokenizer = AutoTokenizer.from_pretrained(model_name)
25
- model = AutoModelForCausalLM.from_pretrained(
26
- model_name,
27
- device_map="auto",
28
- torch_dtype=torch.bfloat16,
29
- attn_implementation="flash_attention_2"
30
- )
31
- return model, tokenizer
32
- except Exception as e:
33
- print(f"ERROR: Failed to load model/tokenizer: {e}")
34
- raise # Re-raise for Gradio
35
-
36
- # --- 3. Chat Template Function ---
37
-
38
- def apply_chat_template(messages, tokenizer):
39
- """Applies the chat template."""
40
- try:
41
- if hasattr(tokenizer, "chat_template") and tokenizer.chat_template:
42
- return tokenizer.apply_chat_template(
43
- messages, tokenize=False, add_generation_prompt=True
44
- )
45
- else:
46
- print("WARNING: Tokenizer lacks chat_template. Using fallback.")
47
- chat_template = "{% for message in messages %}" \
48
- "{{ '<start_of_turn>' + message['role'] + '\n' + message['content'] + '<end_of_turn>\n' }}" \
49
- "{% endfor %}" \
50
- "{% if add_generation_prompt %}{{ '<start_of_turn>model\n' }}{% endif %}"
51
- return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, chat_template=chat_template)
52
- except Exception as e:
53
- print(f"ERROR: Chat template application failed: {e}")
54
- raise
55
-
56
- # --- 4. Text Generation Function ---
57
-
58
- def generate_response(messages, model, tokenizer, max_new_tokens=256, temperature=0.7, top_k=50, top_p=0.95, repetition_penalty=1.2):
59
- """Generates a response."""
60
- prompt = apply_chat_template(messages, tokenizer)
61
- try:
62
- pipeline_instance = pipeline(
63
- "text-generation", model=model, tokenizer=tokenizer,
64
- torch_dtype=torch.bfloat16, device_map="auto",
65
- model_kwargs={"attn_implementation": "flash_attention_2"}
66
- )
67
- outputs = pipeline_instance(
68
- prompt, max_new_tokens=max_new_tokens, do_sample=True,
69
- temperature=temperature, top_k=top_k, top_p=top_p,
70
- repetition_penalty=repetition_penalty, pad_token_id=tokenizer.eos_token_id
71
- )
72
- return outputs[0]["generated_text"][len(prompt):].strip()
73
- except Exception as e:
74
- print(f"ERROR: Response generation failed: {e}")
75
- raise
76
-
77
- # --- 5. Gradio Interface ---
78
- model = None # Initialize model and tokenizer as global variables
79
- tokenizer = None
80
-
81
- def chat(token, message, history):
82
- global model, tokenizer # Access the global model and tokenizer
83
-
84
- if not authenticate(token):
85
- return "Authentication failed. Please enter a valid Hugging Face token.", history
86
-
87
- if model is None or tokenizer is None:
88
- try:
89
- model, tokenizer = load_model_and_tokenizer()
90
- except Exception as e:
91
- return f"Model loading error: {e}", history
92
-
93
- if not history:
94
- history = []
95
- messages = [{"role": "user", "content": msg} for msg, _ in history]
96
- messages.extend([{"role": "model", "content": resp} for _, resp in history if resp])
97
  messages.append({"role": "user", "content": message})
98
 
99
- try:
100
- response = generate_response(messages, model, tokenizer)
101
- history.append((message, response))
102
- return "", history
103
- except Exception as e:
104
- return f"Error during generation: {e}", history
105
-
106
- with gr.Blocks() as demo:
107
- gr.Markdown("# Gemma Chatbot")
108
- gr.Markdown("Enter your Hugging Face API token (read access required):")
109
- token_input = gr.Textbox(label="Hugging Face Token", type="password") # Use type="password"
110
- chatbot = gr.Chatbot(label="Chat", height=400)
111
- msg_input = gr.Textbox(label="Message", placeholder="Ask me anything!")
112
- clear_btn = gr.ClearButton([msg_input, chatbot])
113
-
114
- msg_input.submit(chat, [token_input, msg_input, chatbot], [msg_input, chatbot])
115
- clear_btn.click(lambda: (None, []), [], [msg_input, chatbot])
116
-
117
-
118
- demo.launch()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  import gradio as gr
2
+ from huggingface_hub import InferenceClient
3
+
4
+ """
5
+ For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
6
+ """
7
+ client = InferenceClient("google/gemma-3-1b-it")
8
+
9
+
10
+ def respond(
11
+ message,
12
+ history: list[tuple[str, str]],
13
+ system_message,
14
+ max_tokens,
15
+ temperature,
16
+ top_p,
17
+ ):
18
+ messages = [{"role": "system", "content": system_message}]
19
+
20
+ for val in history:
21
+ if val[0]:
22
+ messages.append({"role": "user", "content": val[0]})
23
+ if val[1]:
24
+ messages.append({"role": "assistant", "content": val[1]})
25
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
26
  messages.append({"role": "user", "content": message})
27
 
28
+ response = ""
29
+
30
+ for message in client.chat_completion(
31
+ messages,
32
+ max_tokens=max_tokens,
33
+ stream=True,
34
+ temperature=temperature,
35
+ top_p=top_p,
36
+ ):
37
+ token = message.choices[0].delta.content
38
+
39
+ response += token
40
+ yield response
41
+
42
+
43
+ """
44
+ For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
45
+ """
46
+ demo = gr.ChatInterface(
47
+ respond,
48
+ additional_inputs=[
49
+ gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
50
+ gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
51
+ gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
52
+ gr.Slider(
53
+ minimum=0.1,
54
+ maximum=1.0,
55
+ value=0.95,
56
+ step=0.05,
57
+ label="Top-p (nucleus sampling)",
58
+ ),
59
+ ],
60
+ )
61
+
62
+
63
+ if __name__ == "__main__":
64
+ demo.launch()