Update README.md

README.md

@@ -53,83 +53,38 @@ If you're eager to have a conversation with VergilGPT2, you can utilize the foll
 
 ```python
 import torch
-import re
-import random
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-def filter_english_words(text):
-    # Initialize an empty list to store the filtered words
-    filtered_text = []
-    # Split the input text into individual words
-    words = text.split()
-    # Iterate through the words
-    for word in words:
-        # Check if the word is a valid English word or a punctuation
-        if word.isalpha() or re.match(r'^[!"#$%&\'()*+,-./:;<=>?@[\]^_`{|}~]$', word):
-            # If it is, append it to the filtered text list
-            filtered_text.append(word)
-    # Return the filtered text as a string
-    return ' '.join(filtered_text)
-
-def generate_response(model, tokenizer, input_text, max_length=300, min_length=20, num_return_sequences=1, temperature=0.2, top_k=50, top_p=0.9, num_beams=10, repetition_penalty=1.0):
-    input_ids = tokenizer.encode(input_text, return_tensors='pt')
-
-    # Set model to eval mode
-    model.eval()
-
-    # Generate responses using different decoding strategies
-    try:
-        output = model.generate(
-            input_ids=input_ids,
-            max_length=max_length,
-            min_length=min_length,
-            num_return_sequences=num_return_sequences,
-            temperature=temperature,
-            top_k=top_k,
-            top_p=top_p,
-            num_beams=num_beams,
-            no_repeat_ngram_size=2,
-            do_sample=True,
-        )
-
-        # Decode the generated responses
-        responses = [tokenizer.decode(o, skip_special_tokens=True) for o in output]
-        # Remove input_text from the responses
-        responses = [response[len(tokenizer.decode(input_ids[0], skip_special_tokens=True)):] for response in responses]
-
-        # Filter out non-English words
-        responses = [filter_english_words(response) for response in responses]
-
-        return responses
-
-    except RuntimeError:
-        return ["I'm sorry, I encountered an error while generating a response."]
-
-
 # Load pre-trained model and tokenizer
 access_token = "REPLACE_WITH_ACCESS_TOKEN"
-model_id = "Starcodium/VergilGPT2"
-tokenizer = AutoTokenizer.from_pretrained(model_id, revision="main", use_auth_token=access_token)
-model = AutoModelForCausalLM.from_pretrained(model_id, revision="main", use_auth_token=access_token)
+model_id = "Starcodium/Vergil_GPT-2"
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=access_token)
+model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=access_token)
 
 tokenizer.pad_token = tokenizer.eos_token
 model.config.pad_token_id = model.config.eos_token_id
 
 # Get user input and generate responses
 while True:
-    input_text = input("Type 'quit' or 'exit' to stop runtime.\nEnter your input text: ")
-    if input_text.lower() in ["quit", "exit"]:
+    # Get user input
+    prompt = input("\nEnter your prompt (or 'exit' to quit): ")
+    if prompt.lower() == 'exit':
         break
 
-    responses = generate_response(model, tokenizer, input_text, max_length=100, min_length=20, num_return_sequences=1, temperature=0.7, top_k=40, top_p=0.5, num_beams=1)
-    responses = [r for r in responses if r.strip() != '']
-    if responses:
-        response = responses[0]
-    else:
-        response = "I'm sorry, I don't have a response at the moment."
+    prompt_template = f"""A chat between a curious user and an artificial intelligence assistant named 'Vergil'. Vergil gives helpful, detailed, and polite answers to the user's questions.
+
+    USER: {prompt}
+    VERGIL:
+    """
+
+    print("\n\nGenerating")
+
+    input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids.cuda()
+    output = model.generate(inputs=input_ids, temperature=0.7, max_new_tokens=512)
+    response = tokenizer.decode(output[0]).replace(prompt_template,"").replace("<s> ", "").replace("</s>", "").split("VERGIL: ")[-1].strip()  # Only keep the model's response
 
-    # Print the bot's response
-    print("Vergil: "+response)
+    # Print only the model's response, without the conversation history
+    print(response)
 ```
 
 This code snippet allows you to engage in conversations with VergilGPT2. Simply enter your input text, and VergilGPT2 will generate responses based on the provided context. Experiment with different values of the variables temperature, top_k, and top_p to customize the response generation process according to your desired preferences.