Views
No views yet
IMPORTANT: This model is sensitive to the chat template used. Ensure you use the correct template:
<s>system
[System message]</s>
<s>user
[Your question or message]</s>
<s>assistant
[The model's response]</s>1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4# Determine the device to use (GPU if available, otherwise CPU)
5device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
6
7# Load the model and tokenizer, then move the model to the appropriate device
8model = AutoModelForCausalLM.from_pretrained("adi2606/MenstrualQA").to(device)
9tokenizer = AutoTokenizer.from_pretrained("adi2606/MenstrualQA")
10
11# Function to generate a response from the chatbot
12def generate_response(message: str, temperature: float = 0.4, repetition_penalty: float = 1.1) -> str:
13 # Apply the chat template and convert to PyTorch tensors
14 messages = [
15 {"role": "system", "content": "You are a helpful assistant."},
16 {"role": "user", "content": message}
17 ]
18 input_ids = tokenizer.apply_chat_template(
19 messages, add_generation_prompt=True, return_tensors="pt"
20 ).to(device)
21
22 # Generate the response
23 output = model.generate(
24 input_ids,
25 max_length=512,
26 temperature=temperature,
27 repetition_penalty=repetition_penalty,
28 do_sample=True
29 )
30
31 # Decode the generated output
32 generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
33 return generated_text
34
35# Example usage
36message = "how to stop pain during menstruation?"
37response = generate_response(message)
38print(response)
39