File size: 1,282 Bytes
e1348f5
4450af9
e1348f5
4827221
 
4450af9
eb6299e
 
4450af9
bf82c59
eb6299e
 
 
a7322e4
eb6299e
bf82c59
 
 
 
 
a7322e4
eb6299e
 
 
a7322e4
bf82c59
a7322e4
eb6299e
 
bf82c59
 
 
eb6299e
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load model directly
tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2")

# System message
system_message = "You are a code teaching assistant named OmniCode created by Anusha K. Answer all the code related questions being asked."


def generate_response(prompt, max_length=150, temperature=1.0):
    input_text = system_message + "\n" + prompt
    input_ids = tokenizer.encode(input_text, return_tensors='pt')

    # Generate response
    output = model.generate(input_ids,
                             max_length=max_length,
                             temperature=temperature,
                             pad_token_id=tokenizer.eos_token_id,
                             num_return_sequences=1)

    # Decode and return the response
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response


if __name__ == "__main__":
    while True:
        user_input = input("You: ")
        if not user_input:  # Check if user input is empty
            print("Exiting OmniCode. Thank you for using me!")
            break
        response = generate_response(user_input)
        print("OmniCode:", response)