Update app.py

app.py

@@ -1,134 +1,101 @@
 import gradio as gr
-import re
-import torch
 import requests
-import random
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
+import re
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 
-# TMDb API key
 TMDB_API_KEY = "364ab334807e87805d326b6dc0f1905c"
-
-# Mapeamento dos gêneros TMDb
 GENRE_MAP = {
-    "Action": 28, "Adventure": 12, "Animation": 16, "Comedy": 35, "Crime": 80,
-    "Documentary": 99, "Drama": 18, "Family": 10751, "Fantasy": 14, "History": 36,
-    "Horror": 27, "Music": 10402, "Mystery": 9648, "Romance": 10749, "Sci-Fi": 878,
-    "TV Movie": 10770, "Thriller": 53, "War": 10752, "Western": 37
+    "Action": 28, "Adventure": 12, "Animation": 16, "Comedy": 35,
+    "Crime": 80, "Documentary": 99, "Drama": 18, "Family": 10751,
+    "Fantasy": 14, "History": 36, "Horror": 27, "Music": 10402,
+    "Mystery": 9648, "Romance": 10749, "Sci-Fi": 878, "TV Movie": 10770,
+    "Thriller": 53, "War": 10752, "Western": 37
 }
 
-# Modelos
+# Modelos pequenos para Hugging Face
+gpt_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
+gpt_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
+gpt_pipe = pipeline("text2text-generation", model=gpt_model, tokenizer=gpt_tokenizer)
 
-gpt4_tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-rw-1b")
-gpt4_model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-rw-1b")
-gpt4_pipe = pipeline("text-generation", model=gpt4_model, tokenizer=gpt4_tokenizer, device=0 if torch.cuda.is_available() else -1)
-
-claude_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
-claude_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
+claude_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
+claude_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
 claude_pipe = pipeline("text2text-generation", model=claude_model, tokenizer=claude_tokenizer)
 
-arb_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
-arb_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
-arb_pipe = pipeline("text2text-generation", model=arb_model, tokenizer=arb_tokenizer)
-
-# Extrai gêneros reconhecidos
-GENRE_LIST = list(GENRE_MAP.keys())
-GENRE_PATTERN = r"\\b(" + "|".join(GENRE_LIST) + r")\\b"
+arb_pipe = claude_pipe  # usa o mesmo modelo para o juiz
 
 def extract_genres(output):
-    genres = re.findall(GENRE_PATTERN, output, flags=re.IGNORECASE)
-    return sorted(set([g.title() for g in genres if g.title() in GENRE_MAP]))
+    allowed_genres = list(GENRE_MAP.keys())
+    return sorted(set(re.findall(r'\b(?:' + '|'.join(allowed_genres) + r')\b', output, flags=re.IGNORECASE)))
 
-# Classificadores de gênero
-
-def gpt4_genre_classifier(text):
+def classify_genres(pipe, text):
     prompt = (
-        "List only the relevant genres present in the following game or movie description. Do not list all possible genres.\n"
+        "List the genres (only) for this movie or game description.\n"
+        "Valid genres: " + ", ".join(GENRE_MAP.keys()) + ".\n"
         f"Description: \"{text}\"\nGenres:"
     )
-    output = gpt4_pipe(prompt, max_new_tokens=20)[0]['generated_text']
-    genres = extract_genres(output)
-    return genres
-
-def claude_genre_classifier(text):
-    prompt = (
-        "List only the relevant genres present in the following game or movie description. Do not list all possible genres.\n"
-        f"Description: \"{text}\"\nGenres:"
-    )
-    output = claude_pipe(prompt, max_new_tokens=20)[0]['generated_text']
-    genres = extract_genres(output)
-    return genres
-
-# Busca um filme aleatório com ao menos 1 gênero identificado
-
-def recommend_movie(genres, used_ids):
-    random.shuffle(genres)
-    for genre in genres:
-        genre_id = GENRE_MAP.get(genre)
-        if not genre_id:
-            continue
-        url = f"https://api.themoviedb.org/3/discover/movie?api_key={TMDB_API_KEY}&with_genres={genre_id}&sort_by=vote_average.desc&vote_count.gte=100"
-        resp = requests.get(url)
-        if resp.status_code == 200:
-            results = resp.json().get("results", [])
-            for movie in results:
-                if movie['id'] not in used_ids:
-                    used_ids.add(movie['id'])
-                    return {
-                        "title": movie['title'],
-                        "rating": movie['vote_average'],
-                        "genre": genre
-                    }
+    output = pipe(prompt, max_new_tokens=30)[0]["generated_text"]
+    return extract_genres(output)
+
+def get_movie_by_genres(genres, exclude_titles=[]):
+    genre_ids = [GENRE_MAP[g] for g in genres if g in GENRE_MAP]
+    if not genre_ids:
+        return None
+    url = f"https://api.themoviedb.org/3/discover/movie"
+    params = {
+        "api_key": TMDB_API_KEY,
+        "with_genres": ",".join(map(str, genre_ids)),
+        "sort_by": "vote_average.desc",
+        "vote_count.gte": 50,
+        "language": "en-US"
+    }
+    response = requests.get(url, params=params).json()
+    movies = response.get("results", [])
+    for movie in movies:
+        if movie["title"] not in exclude_titles:
+            return {"title": movie["title"], "rating": movie["vote_average"]}
     return None
 
-# Árbitro
-
-def llama_judge(text, gpt_genres, claude_genres):
+def judge_and_decide(text, gpt_genres, claude_genres, gpt_movie, claude_movie):
     prompt = (
-        f"You are an expert genre classifier.\n"
-        f"Description: \"{text}\"\n"
-        f"GPT-4 genres: {', '.join(gpt_genres)}\n"
-        f"Claude genres: {', '.join(claude_genres)}\n"
-        "Who identified the genres more accurately? Reply with only 'GPT-4', 'Claude', or 'Tie'."
+        f"You are a genre classification judge.\n"
+        f"Description: {text}\n"
+        f"GPT Genres: {', '.join(gpt_genres)}\nClaude Genres: {', '.join(claude_genres)}\n"
+        f"Which list better matches the description? Reply with GPT, Claude or Tie."
     )
-    output = arb_pipe(prompt, max_new_tokens=10)[0]['generated_text']
-    if "gpt" in output.lower():
-        return "GPT-4"
-    elif "claude" in output.lower():
-        return "Claude"
-    return "Tie"
-
-# Função principal
-
-def process_input(text):
-    gpt_genres = gpt4_genre_classifier(text)
-    claude_genres = claude_genre_classifier(text)
-
-    used_ids = set()
-    gpt_movie = recommend_movie(gpt_genres, used_ids)
-    claude_movie = recommend_movie(claude_genres, used_ids)
+    decision = arb_pipe(prompt, max_new_tokens=10)[0]["generated_text"]
+    decision = "Claude" if "Claude" in decision else "GPT" if "GPT" in decision else "Tie"
+    
+    best_movie = max([gpt_movie, claude_movie], key=lambda m: m["rating"])
+    return {
+        "decision": decision,
+        "winner_movie": best_movie
+    }
+
+def main(description):
+    gpt_genres = classify_genres(gpt_pipe, description)
+    claude_genres = classify_genres(claude_pipe, description)
+
+    gpt_movie = get_movie_by_genres(gpt_genres)
+    claude_movie = get_movie_by_genres(claude_genres, exclude_titles=[gpt_movie["title"]] if gpt_movie else [])
 
     if not gpt_movie or not claude_movie:
         return "⚠️ Não foi possível encontrar recomendações de filmes para os gêneros identificados."
 
-    winner = llama_judge(text, gpt_genres, claude_genres)
-    best_movie = gpt_movie if gpt_movie['rating'] >= claude_movie['rating'] else claude_movie
+    result = judge_and_decide(description, gpt_genres, claude_genres, gpt_movie, claude_movie)
 
     return (
-        f"GPT-4 Genres: {', '.join(gpt_genres)}\n"
-        f"Claude Genres: {', '.join(claude_genres)}\n\n"
-        f"GPT-4 Movie: {gpt_movie['title']} ({gpt_movie['rating']})\n"
-        f"Claude Movie: {claude_movie['title']} ({claude_movie['rating']})\n\n"
-        f"🏆 Winner: {winner}\n"
-        f"🎬 Best Rated Movie: {best_movie['title']}"
+        f"🎯 GPT-4 Genres: {', '.join(gpt_genres)}\n"
+        f"🎯 Claude Genres: {', '.join(claude_genres)}\n\n"
+        f"🎬 GPT Movie: {gpt_movie['title']} ({gpt_movie['rating']})\n"
+        f"🎬 Claude Movie: {claude_movie['title']} ({claude_movie['rating']})\n\n"
+        f"🏆 Winner: {result['decision']}\n"
+        f"🎖️ Best Rated Movie: {result['winner_movie']['title']} ({result['winner_movie']['rating']})"
     )
 
-# Interface Gradio
-iface = gr.Interface(
-    fn=process_input,
-    inputs=gr.Textbox(lines=3, label="Enter game or movie description in English"),
+gr.Interface(
+    fn=main,
+    inputs=gr.Textbox(lines=3, placeholder="Enter a movie/game description in English"),
     outputs="text",
-    title="AI Genre Duel + Movie Recommendation",
-    description="Dois modelos de IA identificam os gêneros de uma descrição e recomendam um filme com base nisso. Um juiz decide o melhor classificador e destaca o filme com maior nota."
-)
-
-iface.launch(share=True)
+    title="🎬 AI Genre Duel + Movie Recommendation",
+    description="Dois modelos AI classificam o gênero de uma descrição e recomendam filmes. O juiz escolhe o vencedor com base nos gêneros e na melhor nota de filme."
+).launch()