Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2from peft import PeftModel, PeftConfig
3from huggingface_hub import login
4
5# load model for probability predicting
6peft_model_id = "mariannam/llammas-prediction-grading"
7config = PeftConfig.from_pretrained(peft_model_id)
8
9# load base model
10base_model = AutoModelForCausalLM.from_pretrained(
11 config.base_model_name_or_path,
12 device_map="auto", # for CPU use device_map=None
13 torch_dtype="auto"
14)
15
16# load adapters
17model = PeftModel.from_pretrained(base_model, peft_model_id)
18
19# load tokenizer
20tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
21
22# input ocr text HERE
23ocr_text = "Misso wallas awati awalik telefoni-kõnc-punkt Hinol ja PiiganbiS mõlemal, Kanepi kanbu."
24
25# input corrected text you want the probability for HERE
26prediction = "Misso wallas awati awalik telefoni-kõne-punkt Hinol ja Piigandil mõlemal, Kanepi kandu."
27
28# prompt tempelate
29prompt = f"""### Instruction:
30Kui suur on tõenäosus, et parandatud tekst on OCR tekstist parem? Tagasta tõenäosus täisarvulise protsendina.
31
32### Input:
33OCR TEKST: {ocr_text}
34
35PARANDATUD TEKST: {prediction}
36
37### Response:
38"""
39
40# generate and print output
41inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
42outputs = model.generate(**inputs, max_new_tokens=3)
43print(tokenizer.decode(outputs[0], skip_special_tokens=True)[len(prompt):].strip())