Dhahlan2000 commited on
Commit
825d8ee
·
verified ·
1 Parent(s): 9b3085b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +49 -0
app.py ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from huggingface_hub import InferenceClient
3
+ from transformers import AutoTokenizer, AutoModelForCausalLM
4
+ import torch
5
+ import os
6
+
7
+ # Replace 'your_huggingface_token' with your actual Hugging Face access token
8
+ access_token = os.getenv('token')
9
+
10
+ # Initialize the tokenizer and model with the Hugging Face access token
11
+ tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it", use_auth_token=access_token)
12
+ model = AutoModelForCausalLM.from_pretrained(
13
+ "google/gemma-2b-it",
14
+ torch_dtype=torch.bfloat16,
15
+ use_auth_token=access_token
16
+ )
17
+ model.eval() # Set the model to evaluation mode
18
+
19
+ # Initialize the inference client (if needed for other API-based tasks)
20
+ client = InferenceClient(token=access_token)
21
+
22
+ def conversation_predict(input_text):
23
+ """Generate a response for single-turn input using the model."""
24
+ # Tokenize the input text
25
+ input_ids = tokenizer(input_text, return_tensors="pt").input_ids
26
+
27
+ # Generate a response with the model
28
+ outputs = model.generate(input_ids, max_new_tokens=2048)
29
+
30
+ # Decode and return the generated response
31
+ return tokenizer.decode(outputs[0], skip_special_tokens=True)
32
+
33
+ def respond():
34
+ """Streamlit app for a multi-turn chat conversation."""
35
+ st.title("Chat with Gemma")
36
+
37
+ system_message = st.text_input("System message", value="You are a friendly Chatbot.")
38
+ max_tokens = st.slider("Max new tokens", min_value=1, max_value=2048, value=512, step=1)
39
+ temperature = st.slider("Temperature", min_value=0.1, max_value=4.0, value=0.7, step=0.1)
40
+ top_p = st.slider("Top-p (nucleus sampling)", min_value=0.1, max_value=1.0, value=0.95, step=0.05)
41
+
42
+ message = st.text_input("Your message")
43
+
44
+ if message:
45
+ response = conversation_predict(message)
46
+ st.write(response)
47
+
48
+ if __name__ == "__main__":
49
+ respond()