from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import gradio as gr

# Download the model from Hugging Face
model_name = "johnpaulbin/articulate-V1-Q8_0-GGUF"
model_file = "articulate-V1-q8_0.gguf"  # Verify the exact file name in the repository
model_path = hf_hub_download(repo_id=model_name, filename=model_file)

# Initialize the Llama model
llm = Llama(
    model_path=model_path,
    n_ctx=1028,      # Context length
    n_threads=2,    # Number of CPU threads
    n_gpu_layers=0   # Run on CPU only
)

# Define the translation function
def translate(direction, text):
    # Determine source and target languages based on direction
    if direction == "English to Spanish":
        source_lang = "ENGLISH"
        target_lang = "SPANISH"
    elif direction == "Spanish to English":
        source_lang = "SPANISH"
        target_lang = "ENGLISH"
    else:
        return "Invalid direction"
    
    # Construct the prompt for raw completion
    prompt = f"[{source_lang}]{text}[{target_lang}]"
    
    # Generate completion with deterministic settings (greedy decoding)
    response = llm.create_completion(
        prompt,
        max_tokens=200,    # Limit output length
        temperature=0,     # Greedy decoding
        top_k=1            # Select the most probable token
    )
    
    # Extract and return the generated text
    return response['choices'][0]['text'].strip()

# Define the Gradio interface
direction_options = ["English to Spanish", "Spanish to English"]
iface = gr.Interface(
    fn=translate,
    inputs=[
        gr.Dropdown(choices=direction_options, label="Translation Direction"),
        gr.Textbox(lines=5, label="Input Text")
    ],
    outputs=gr.Textbox(lines=5, label="Translation"),
    title="Translation App",
    description="Translate text between English and Spanish using the Articulate V1 model."
)

# Launch the app
iface.launch()