Views
No views yet
pip3 install transformers1from typing import List, Dict
2import torch
3from transformers import pipeline
4
5def generate_prompts_for_classification(article: str, summary_sentences: List[str]) -> List[Dict]:
6 prompts = []
7 for sentence in summary_sentences:
8 prompt = {"text": article, "text_pair": sentence}
9 prompts.append(prompt)
10 return prompts
11
12def predict_with_hf_classification_pipeline(prompts: List[Dict], model_name: str, max_context_length: int = 512,
13 batch_size: int = 2) -> List[str]:
14 device = "cuda" if torch.cuda.is_available() else "cpu"
15 text_classification_pipeline = pipeline("text-classification", model=model_name, device=device,
16 batch_size=batch_size)
17
18 batch_output = text_classification_pipeline(prompts, truncation=True, max_length=max_context_length)
19 predictions = [result['label'] for result in batch_output]
20 return predictions
21
22def main():
23
24 model_name = "mtc/mbert-absinth-3-epochs"
25 # Articles longer than 512 tokens will be truncated
26 max_context_length = 512
27 # Adjust batch_size according to your local gpu memory
28 batch_size = 2
29
30 article = "Ein neuer Zirkus ist gestern in Zürich angekommen. Viele Familien besuchten das grosse Zelt, um die Vorstellung zu sehen. Es gab Akrobaten, Clowns und Tiere, die das Publikum begeisterten. Der Zirkus bleibt noch eine Woche in der Stadt und bietet täglich Vorstellungen an."
31
32 summary_sentences = [
33 "Ein Zirkus ist in Basel angekommen.",
34 "Der Zirkus, der in 1950 gegründet wurde, wird von vielen Familien besucht."]
35
36 prompts = generate_prompts_for_classification(article=article, summary_sentences=summary_sentences)
37 predictions = predict_with_hf_classification_pipeline(prompts=prompts, model_name=model_name,
38 max_context_length=max_context_length, batch_size=batch_size)
39 print(predictions)
40
41
42if __name__ == '__main__':
43 main()[ Intrinsic Hallucination
Extrinsic Hallucination]