Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -10,24 +10,23 @@ model = AutoModelForCausalLM.from_pretrained("adi2606/MenstrualQA").to(device)
|
|
10 |
tokenizer = AutoTokenizer.from_pretrained("adi2606/MenstrualQA")
|
11 |
|
12 |
# Function to generate a response from the chatbot
|
13 |
-
def generate_response(message: str, temperature: float = 0.4, repetition_penalty: float = 1.1
|
14 |
-
|
15 |
-
|
16 |
-
|
17 |
-
|
18 |
-
|
19 |
-
|
|
|
20 |
).to(device)
|
21 |
|
22 |
# Generate the response
|
23 |
output = model.generate(
|
24 |
-
|
25 |
-
attention_mask=inputs['attention_mask'],
|
26 |
max_length=512,
|
27 |
temperature=temperature,
|
28 |
repetition_penalty=repetition_penalty,
|
29 |
-
do_sample=True
|
30 |
-
pad_token_id=tokenizer.eos_token_id
|
31 |
)
|
32 |
|
33 |
# Decode the generated output
|
|
|
10 |
tokenizer = AutoTokenizer.from_pretrained("adi2606/MenstrualQA")
|
11 |
|
12 |
# Function to generate a response from the chatbot
|
13 |
+
def generate_response(message: str, temperature: float = 0.4, repetition_penalty: float = 1.1) -> str:
|
14 |
+
# Apply the chat template and convert to PyTorch tensors
|
15 |
+
messages = [
|
16 |
+
{"role": "system", "content": "You are a helpful assistant."},
|
17 |
+
{"role": "user", "content": message}
|
18 |
+
]
|
19 |
+
input_ids = tokenizer.apply_chat_template(
|
20 |
+
messages, add_generation_prompt=True, return_tensors="pt"
|
21 |
).to(device)
|
22 |
|
23 |
# Generate the response
|
24 |
output = model.generate(
|
25 |
+
input_ids,
|
|
|
26 |
max_length=512,
|
27 |
temperature=temperature,
|
28 |
repetition_penalty=repetition_penalty,
|
29 |
+
do_sample=True
|
|
|
30 |
)
|
31 |
|
32 |
# Decode the generated output
|