CreitinGameplays commited on
Commit
b793725
1 Parent(s): 33f79d4

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +14 -7
app.py CHANGED
@@ -9,14 +9,17 @@ model_name = "CreitinGameplays/bloom-3b-conversational"
9
  tokenizer = AutoTokenizer.from_pretrained(model_name)
10
  model = AutoModelForCausalLM.from_pretrained(model_name)
11
 
12
- def generate_text(prompt):
13
- """Generates text using the BLOOM model from Hugging Face Transformers."""
14
- # Encode the prompt into tokens
15
- input_ids = tokenizer(prompt, return_tensors="pt").input_ids
16
 
17
- # Generate text with the prompt and limit the maximum length to 256 tokens
 
 
 
18
  output = model.generate(
19
- input_ids=input_ids,
20
  max_length=256,
21
  num_beams=1,
22
  num_return_sequences=1, # Generate only 1 sequence
@@ -29,7 +32,11 @@ def generate_text(prompt):
29
 
30
  # Decode the generated token sequence back to text
31
  generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
32
- return generated_text
 
 
 
 
33
 
34
  # Define the Gradio interface
35
  interface = gr.Interface(
 
9
  tokenizer = AutoTokenizer.from_pretrained(model_name)
10
  model = AutoModelForCausalLM.from_pretrained(model_name)
11
 
12
+ def generate_text(user_prompt):
13
+ """Generates text using the BLOOM model from Hugging Face Transformers and removes the user prompt."""
14
+ # Construct the full prompt with system introduction, user prompt, and assistant role
15
+ prompt = f"<|system|> You are a helpful AI assistant. </s> <|prompter|> {user_prompt} </s> <|assistant|>"
16
 
17
+ # Encode the entire prompt into tokens
18
+ prompt_encoded = tokenizer(prompt, return_tensors="pt").input_ids
19
+
20
+ # Generate text with the complete prompt and limit the maximum length to 256 tokens
21
  output = model.generate(
22
+ input_ids=prompt_encoded,
23
  max_length=256,
24
  num_beams=1,
25
  num_return_sequences=1, # Generate only 1 sequence
 
32
 
33
  # Decode the generated token sequence back to text
34
  generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
35
+
36
+ # Extract the assistant's response (assuming it starts with "<|assistant|>")
37
+ assistant_response = generated_text.split("<|assistant|>")[-1]
38
+
39
+ return assistant_response
40
 
41
  # Define the Gradio interface
42
  interface = gr.Interface(