load model

app.py

@@ -7,6 +7,23 @@ hub_path = 'guptavishal79/aimlops'
 loaded_model = GPT2LMHeadModel.from_pretrained(hub_path)
 loaded_tokenizer = GPT2Tokenizer.from_pretrained(hub_path)
 
+
+def generate_response(model, tokenizer, prompt, max_length=200):
+  input_ids = tokenizer.encode(prompt, return_tensors="pt")
+  # Create the attention mask and pad token id
+  attention_mask = torch.ones_like(input_ids)
+  pad_token_id = tokenizer.eos_token_id
+
+  output = model.generate(
+      input_ids,
+      max_length=max_length,
+      num_return_sequences=1,
+      attention_mask=attention_mask,
+      pad_token_id=pad_token_id
+  )
+
+  return tokenizer.decode(output[0], skip_special_tokens=True)
+
 # Function for response generation
 def generate_query_response(prompt, max_length=200):