Update app.py

app.py

@@ -1,144 +1,98 @@
 import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModel
+from transformers import pipeline
 import torch
 import warnings
 warnings.filterwarnings('ignore')
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-try:
-    models = {
-        'transcription': pipeline("automatic-speech-recognition", 
-                                model="openai/whisper-small", 
-                                device=device),
-        'translation': pipeline("translation", 
-                              model="facebook/mbart-large-50-many-to-many-mmt",
-                              device=device),
-        'summarization': pipeline("summarization", 
-                                model="facebook/bart-large-cnn", 
-                                device=device),
-        'sentiment': pipeline("sentiment-analysis", 
-                            model="nlptown/bert-base-multilingual-uncased-sentiment", 
-                            device=device),
-        'question_answering': pipeline("question-answering", 
-                                     model="deepset/roberta-base-squad2", 
-                                     device=device),
-        'chat': pipeline("text-generation", 
-                        model="facebook/opt-125m", 
-                        device=device)
-    }
-    
-except Exception as e:
-    print(f"Erro ao carregar modelos: {str(e)}")
-
-def safe_process(func):
-    def wrapper(*args, **kwargs):
-        try:
-            return func(*args, **kwargs)
-        except Exception as e:
-            return f"Erro ao processar: {str(e)}"
-    return wrapper
-
-@safe_process
-def transcribe(audio):
-    if not audio:
-        return "Por favor, forneça um arquivo de áudio."
-    return models['transcription'](audio)["text"]
-
-@safe_process
-def translate(text, direction):
-    if not text:
-        return "Por favor, insira um texto para tradução."
-    
-    if direction == "pt_en":
-        result = models['translation'](text, src_lang="pt", tgt_lang="en")[0]
-    else:
-        result = models['translation'](text, src_lang="en", tgt_lang="pt")[0]
-    
-    return result['translation_text']
-
-@safe_process
-def summarize(text):
-    if not text:
-        return "Por favor, insira um texto para resumir."
-    return models['summarization'](text, max_length=130, min_length=30)[0]['summary_text']
-
-@safe_process
-def analyze_sentiment(text):
-    if not text:
-        return "Por favor, insira um texto para análise."
-    return models['sentiment'](text)[0]['label']
-
-@safe_process
-def answer_question(question, context):
-    if not question or not context:
-        return "Por favor, forneça tanto a pergunta quanto o contexto."
-    return models['question_answering'](question=question, context=context)['answer']
+# Expanded models dictionary with more services
+models = {
+    'transcription': pipeline("automatic-speech-recognition", model="openai/whisper-small", device=device),
+    'translation': pipeline("translation", model="facebook/mbart-large-50-many-to-many-mmt", device=device),
+    'summarization': pipeline("summarization", model="facebook/bart-large-cnn", device=device),
+    'sentiment': pipeline("sentiment-analysis", model="nlptown/bert-base-multilingual-uncased-sentiment", device=device),
+    'question_answering': pipeline("question-answering", model="deepset/roberta-base-squad2", device=device),
+    'chat': pipeline("text-generation", model="facebook/opt-125m", device=device),
+    'image_caption': pipeline("image-to-text", model="Salesforce/blip-image-captioning-base", device=device),
+    'text_to_speech': pipeline("text-to-audio", model="facebook/mms-tts-eng", device=device),
+    'zero_shot': pipeline("zero-shot-classification", model="facebook/bart-large-mnli", device=device),
+    'ner': pipeline("token-classification", model="dslim/bert-base-NER", device=device)
+}
 
-@safe_process
-def chat_response(message, history):
-    if not message:
-        return [], history
-    response = models['chat'](message, max_length=100, do_sample=True, temperature=0.7)
-    history.append((message, response[0]['generated_text']))
-    return "", history
+def process_request(service_type, *args, **kwargs):
+    try:
+        if not args[0]: return "Por favor, forneça os dados necessários."
+        
+        processors = {
+            'transcription': lambda audio: models['transcription'](audio)["text"],
+            'translation': lambda text, direction: models['translation'](text, 
+                src_lang="pt" if direction=="pt_en" else "en", 
+                tgt_lang="en" if direction=="pt_en" else "pt")[0]['translation_text'],
+            'summarization': lambda text: models['summarization'](text, max_length=130, min_length=30)[0]['summary_text'],
+            'sentiment': lambda text: models['sentiment'](text)[0]['label'],
+            'qa': lambda q, ctx: models['question_answering'](question=q, context=ctx)['answer'],
+            'chat': lambda msg, history: (models['chat'](msg, max_length=100, do_sample=True)[0]['generated_text'], history + [(msg, "")]),
+            'image_caption': lambda img: models['image_caption'](img)[0]['generated_text'],
+            'tts': lambda text: models['text_to_speech'](text),
+            'zero_shot': lambda text, labels: models['zero_shot'](text, labels)['labels'][0],
+            'ner': lambda text: models['ner'](text)
+        }
+        
+        return processors[service_type](*args)
+    except Exception as e:
+        return f"Erro: {str(e)}"
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    with gr.Tab("Início"):
-        gr.HTML(open("index.html").read())
-    
-    with gr.Tab("Transcrição de Áudio"):
-        audio_input = gr.Audio(type="filepath", label="Upload de Áudio")
-        transcribe_button = gr.Button("Transcrever")
-        transcription_output = gr.Textbox(label="Transcrição", lines=3)
-        transcribe_button.click(transcribe, inputs=audio_input, outputs=transcription_output)
-    
-    with gr.Tab("Tradução"):
-        with gr.Row():
-            translation_direction = gr.Radio(
-                ["en_pt", "pt_en"], 
-                label="Direção da Tradução",
-                value="en_pt"
-            )
-        text_to_translate = gr.Textbox(label="Texto para Traduzir", lines=3)
-        translate_button = gr.Button("Traduzir")
-        translation_output = gr.Textbox(label="Tradução", lines=3)
-        translate_button.click(
-            translate, 
-            inputs=[text_to_translate, translation_direction], 
-            outputs=translation_output
-        )
-    
-    with gr.Tab("Resumo"):
-        text_to_summarize = gr.Textbox(label="Texto para Resumir", lines=5)
-        summarize_button = gr.Button("Resumir")
-        summary_output = gr.Textbox(label="Resumo", lines=3)
-        summarize_button.click(summarize, inputs=text_to_summarize, outputs=summary_output)
+    gr.HTML(open("index.html").read())
     
-    with gr.Tab("Análise de Sentimento"):
-        sentiment_text = gr.Textbox(label="Texto para Análise", lines=3)
-        sentiment_button = gr.Button("Analisar")
-        sentiment_output = gr.Textbox(label="Sentimento")
-        sentiment_button.click(analyze_sentiment, inputs=sentiment_text, outputs=sentiment_output)
-    
-    with gr.Tab("Perguntas e Respostas"):
-        question_input = gr.Textbox(label="Pergunta")
-        context_input = gr.Textbox(label="Contexto", lines=5)
-        qa_button = gr.Button("Responder")
-        qa_output = gr.Textbox(label="Resposta", lines=2)
-        qa_button.click(
-            answer_question, 
-            inputs=[question_input, context_input], 
-            outputs=qa_output
-        )
-    
-    with gr.Tab("Chat"):
-        chatbot = gr.Chatbot()
-        msg = gr.Textbox(label="Mensagem")
-        clear = gr.Button("Limpar")
+    with gr.Tabs():
+        with gr.TabItem("📝 Transcrição"):
+            audio_input = gr.Audio(type="filepath")
+            transcribe_button = gr.Button("Transcrever")
+            transcription_output = gr.Textbox(label="Resultado")
+            transcribe_button.click(lambda x: process_request('transcription', x), inputs=audio_input, outputs=transcription_output)
+        
+        with gr.TabItem("🔄 Tradução"):
+            with gr.Row():
+                text_to_translate = gr.Textbox(label="Texto")
+                translation_direction = gr.Radio(["en_pt", "pt_en"], value="en_pt")
+            translate_button = gr.Button("Traduzir")
+            translation_output = gr.Textbox(label="Resultado")
+            translate_button.click(lambda x, y: process_request('translation', x, y), 
+                                inputs=[text_to_translate, translation_direction], 
+                                outputs=translation_output)
+        
+        with gr.TabItem("📊 Análise"):
+            with gr.Tabs():
+                with gr.TabItem("Resumo"):
+                    text_sum = gr.Textbox(label="Texto", lines=5)
+                    sum_button = gr.Button("Resumir")
+                    sum_output = gr.Textbox(label="Resultado")
+                    sum_button.click(lambda x: process_request('summarization', x), inputs=text_sum, outputs=sum_output)
+                
+                with gr.TabItem("Sentimento"):
+                    text_sent = gr.Textbox(label="Texto")
+                    sent_button = gr.Button("Analisar")
+                    sent_output = gr.Textbox(label="Resultado")
+                    sent_button.click(lambda x: process_request('sentiment', x), inputs=text_sent, outputs=sent_output)
         
-        msg.submit(chat_response, inputs=[msg, chatbot], outputs=[msg, chatbot])
-        clear.click(lambda: None, None, chatbot, queue=False)
+        with gr.TabItem("🤖 IA Avançada"):
+            with gr.Tabs():
+                with gr.TabItem("Chat"):
+                    chatbot = gr.Chatbot()
+                    msg = gr.Textbox(label="Mensagem")
+                    clear = gr.Button("Limpar")
+                    msg.submit(lambda x, y: process_request('chat', x, y), inputs=[msg, chatbot], outputs=[msg, chatbot])
+                    clear.click(lambda: None, None, chatbot, queue=False)
+                
+                with gr.TabItem("Classificação"):
+                    zero_text = gr.Textbox(label="Texto")
+                    zero_labels = gr.Textbox(label="Categorias (separadas por vírgula)")
+                    zero_button = gr.Button("Classificar")
+                    zero_output = gr.Textbox(label="Resultado")
+                    zero_button.click(lambda x, y: process_request('zero_shot', x, y.split(',')), 
+                                    inputs=[zero_text, zero_labels], 
+                                    outputs=zero_output)
 
-if __name__ == "__main__":
-    demo.launch(share=True)
+demo.launch(share=True)