T5-base-WMT14v2

from transformers import T5Tokenizer, T5ForConditionalGeneration

# Load the pre-trained model and tokenizer
model_name = "SynapseQAI/T5-base-WMT14v2"
tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name)

# Function to translate using a mixed strategy: Beam Search + Top-P Sampling
def translate(sentence, num_beams=4, top_p=0.9):
    # Prepare the input with the prompt for the model
    input_text = f"correct phrasing and grammar:{sentence}"
    input_ids = tokenizer(input_text, return_tensors="pt").input_ids

    # Step 1: Beam search to generate top candidate outputs
    beam_outputs = model.generate(
        input_ids, 
        num_beams=num_beams,  # Define number of beams for beam search
        max_length=40, 
        early_stopping=True,
        num_return_sequences=num_beams,  # Return the top beams for Top-p sampling
        do_sample=False  # Beam search is deterministic
    )

    # Step 2: Apply Top-P sampling on the beam outputs to introduce variety
    top_p_output = model.generate(
        input_ids,
        do_sample=True,
        top_p=top_p,  # Apply Top-p sampling
        max_length=40,
        early_stopping=True
    )

    # Decode the beam outputs and top-p sampled output
    beam_translations = [tokenizer.decode(output, skip_special_tokens=True) for output in beam_outputs]
    top_p_translation = tokenizer.decode(top_p_output[0], skip_special_tokens=True)

    # Clean the translations to remove the prompt text
    clean_top_p_translation = top_p_translation.split(":", 1)[-1].strip()
    
    # Combine top translations (Beam search + Top-P sampling)
    combined_translations = [translation.split(":", 1)[-1].strip() for translation in beam_translations]
    
    # Return both top-p sampled translation and beam search candidates
    return clean_top_p_translation, combined_translations

# French sentences from easy to advanced, including more difficult ones
sentences = [
    "Les scientifiques travaillent dur pour trouver un remède.",
    "La capitale de la France est Paris.",
    "Je voudrais un café s'il vous plaît.",
    "Il fait beau aujourd'hui.",
    "J'aime lire des livres et regarder des films pendant mon temps libre.",
    "Si j'avais su que tu venais, j'aurais préparé quelque chose de spécial pour le dîner.",
    "Même si les avancées technologiques apportent de nombreux avantages, elles posent également des défis éthiques considérables qu'il nous faut relever.",
    "La notion de vérité absolue reste controversée parmi les philosophes depuis l'Antiquité.",
    "Dans une société hyperconnectée, la surabondance d'informations peut rendre difficile la distinction entre les faits et les opinions.",
    "L'impact des changements climatiques sur la biodiversité est d'une ampleur difficile à concevoir.",
    "Il est impératif que nous abordions les crises environnementales avec une approche systémique et collaborative.",
    "Les révolutions industrielles successives ont façonné le monde moderne de manière irréversible.",
    "La théorie des jeux est souvent utilisée pour modéliser des situations où les décisions de plusieurs individus s'influencent mutuellement.",
    "La métamorphose de la chenille en papillon est un processus biologique fascinant et complexe.",
    "Les découvertes en physique quantique ont radicalement transformé notre compréhension de l'univers à l'échelle microscopique.",
    "L'analyse des textes anciens révèle souvent des perspectives historiques inexplorées.",
    "Les neurosciences nous offrent une fenêtre sur le fonctionnement des processus cognitifs et émotionnels du cerveau.",
    "L'interaction entre les cultures a enrichi l'histoire humaine, mais a aussi souvent conduit à des conflits violents.",
    "La nature fractale de certaines structures naturelles remet en question nos notions classiques de géométrie.",
    "Il est étonnant de voir comment la technologie de pointe a rendu possible des prouesses médicales autrefois inimaginables.",
    "La complexité de l'économie mondiale réside dans l'interconnexion des marchés financiers, des politiques monétaires et des échanges internationaux.",
    "L'évolution des espèces s'effectue sur des millions d'années, à travers un processus d'adaptation et de sélection naturelle.",
    "Les concepts d'intelligence artificielle soulèvent des questions éthiques sur l'autonomie et la responsabilité des machines.",
    "Le paradoxe du temps est un thème récurrent dans la science-fiction, souvent utilisé pour explorer des réalités alternatives.",
    "Les travaux de recherche en biologie synthétique promettent de révolutionner la manière dont nous produisons des médicaments.",
    "La théorie des cordes cherche à unifier les forces fondamentales de la nature dans une seule théorie cohérente.",
    "L'éducation joue un rôle central dans la promotion du développement durable et l'atténuation des inégalités mondiales."
]

# Translate each sentence and print the top-p sampled translation + beam search candidates
for sentence in sentences:
    top_p_translation, beam_translations = translate(sentence)
    print(f"French: {sentence}\nTop-P Translation: {top_p_translation}\nBeam Search Candidates: {beam_translations}\n")