Update src/brain.py

src/brain.py

@@ -1,11 +1,11 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-model_name = "meta-llama/Llama-2-7b-hf"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
+tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
+model = GPT2LMHeadModel.from_pretrained("distilgpt2")
 
 def generate_answers(query):
-    input_ids = tokenizer(query, return_tensors="pt").input_ids
-    output = model.generate(input_ids, max_new_tokens=100)
-    return tokenizer.decode(output[0], skip_special_tokens=True)
+    input_ids = tokenizer.encode(input_text, return_tensors='pt')
+    max_length = input_ids.shape[1] + 100
+    generated_ids = model.generate(input_ids, max_length=max_length)
+    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    return generated_text