Update README.md

README.md

@@ -56,31 +56,11 @@ tokenizer = PegasusTokenizer.from_pretrained(model_path)
 model = model.to(device)
 
 
-# Define the inference function
-def inference(input_text):
-    # Tokenize input text
-    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
-    input_ids = inputs["input_ids"].to(device)
-    attention_mask = inputs["attention_mask"].to(device)
-
-    # Generate summary
-    model.eval()
-    with torch.no_grad():
-        outputs = model.generate(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            max_new_tokens=150,
-            num_beams=8,
-            early_stopping=True,
-        )
-
-    # Decode the generated text
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return generated_text
-
-
-# Test with a sample input
-test_input = """
+from transformers import PegasusForConditionalGeneration, PegasusTokenizer
+
+model = PegasusForConditionalGeneration.from_pretrained("Mudasir692/peguses_chat_sum")
+tokenizer = PegasusTokenizer.from_pretrained("Mudasir692/peguses_chat_sum")
+input_text = """
 #Person1#: Hey Alice, congratulations on your promotion!
 #Person2#: Thank you so much! It means a lot to me. I’m still processing it, honestly.
 #Person1#: You totally deserve it. Your hard work finally paid off. Let’s celebrate this weekend.
@@ -90,6 +70,10 @@ test_input = """
 #Person1#: That’s wonderful! I’ll bring a gift too. It’s such a big milestone for you.
 #Person2#: You’re the best. Thanks for always being so supportive.
 """
-print("Summary Text:", inference(test_input))
+inputs = tokenizer(input_text, return_tensors="pt")
+model.eval()
+outputs = model.generate(**inputs, max_new_tokens=100)
+generated_summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
+print("generated summary", generated_summary)