from transformers import pipeline, Gemma2ForSequenceClassification, AutoTokenizer
import numpy as np
import pandas as pd
import torch

model_name = "SinclairSchneider/german_politic_direction_gemma-2-9b"
model = Gemma2ForSequenceClassification.from_pretrained(model_name, dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained(model_name)
pipe = pipeline("text-classification", model=model, tokenizer=tokenizer, top_k=None)

vectors = np.array([[-1, 0], 
                    [-9.99193435e-01,  4.01556900e-02], 
                    [-7.91445449e-01,  6.11239806e-01], 
                    [ 3.82683432e-01,  9.23879533e-01], 
                    [ 8.69790824e-01,  4.93420634e-01], 
                    [1, 0]])

def classify(text):
    classification_result = np.array(pd.DataFrame(pipe(text)[0]).sort_values(by=['label'], key=lambda x: x.map({'DIE LINKE':0, 
                                                                                                                'BÜNDNIS 90/DIE GRÜNEN':1, 
                                                                                                                'SPD':2, 
                                                                                                                'FDP':3, 
                                                                                                                'CDU/CSU':4, 
                                                                                                                'AfD':5}))['score'])
    return float(np.arctan2(*classification_result@vectors)/(np.pi/2))

#Links
print(classify("Wir brauchen eine Vermögensteuer, um den Sozialstaat nachhaltig zu finanzieren."))
#-0.7613435819529438
print(classify("Mietendeckel und mehr gemeinnütziger Wohnungsbau sollen Wohnen bezahlbar machen."))
#-0.747022752207469
print(classify("Die Energiewende muss mit massiven öffentlichen Investitionen beschleunigt werden."))
#-0.7165234574290826


#Mitte
print(classify("Die soziale Marktwirtschaft braucht moderne Regeln und weniger Bürokratie."))
#0.24816468602492553
print(classify("Gezielte Entlastungen für kleine und mittlere Einkommen stärken die Mitte."))
#-0.23390688585648964
print(classify("Bildungsoffensive: Basiskompetenzen sichern, Weiterbildung im Beruf fördern."))
#-0.010101430791014977

#Rechts
print(classify("Deutsche Leitkultur und Sprache stärker in öffentlichen Einrichtungen betonen."))
#0.9658786216889841
print(classify("Grenzschutz an EU-Außengrenzen verstärken, Sekundärmigration begrenzen."))
#0.668343040925338
print(classify("Identitätspolitik an Schulen und Behörden zurückfahren, Fokus auf Leistungsprinzip."))
#0.9935253923542486