#from transformers import AutoModel, AutoTokenizer, pipeline
import shutil
import gradio as gr
import os
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# Încarcă modelul și tokenizer-ul
model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-j-6b")
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6b")


# Locația cache-ului Hugging Face
cache_dir = os.path.join(os.getenv('TRANSFORMERS_CACHE', './cache'))

# Șterge cache-ul dacă dorești să forțezi reîncărcarea modelului
shutil.rmtree(cache_dir, ignore_errors=True)

# Încarcă modelul doar o singură dată, la început
#qa_pipeline = pipeline("question-answering", model="openai-community/gpt2")
# Configurare manuală a pipeline-ului
qa_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Funcția pentru procesarea întrebărilor
def answer_question(context, question):
    # Verificăm că textul de context nu este gol
    if not context.strip():
        return "Te rog să introduci un fragment de text."
    
    # Verificăm că întrebarea nu este goală
    if not question.strip():
        return "Te rog să introduci o întrebare."
    
    result = qa_pipeline({
        "context": context,
        "question": question
    })
    return result['answer']

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Text Fragment Analyzer")
    
    # Input pentru context și întrebare
    context = gr.Textbox(label="Fragmente de text", lines=10, placeholder="Introdu textul aici...")
    question = gr.Textbox(label="Întrebare", placeholder="Ce vrei să întrebi despre text?")
    
    # Output pentru răspuns
    answer = gr.Textbox(label="Răspuns", interactive=False)
    
    # Buton pentru generarea răspunsului
    btn = gr.Button("Generează răspuns")
    
    # Asociem acțiunea butonului cu funcția `answer_question`
    btn.click(answer_question, inputs=[context, question], outputs=answer)

# Rulează aplicația pe Hugging Face
demo.launch(share=True)  # share=True pentru a obține un link public