Update app.py

app.py

@@ -1,46 +1,12 @@
-
 import gradio as gr
 import google.generativeai as genai
-import requests
-import json
 import os
 import time
-
+\
 genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
 gemini_model = genai.GenerativeModel("gemini-2.0-flash")
 
-class OpenRouter:
-    def __init__(self, model_name, key, role="user"):
-        self.model_name = model_name
-        self.key = key
-        self.api_url = "https://openrouter.ai/api/v1/chat/completions"
-        self.role = role
-
-    def get_response(self, prompt):
-        headers = {
-            "Authorization": f"Bearer {self.key}",
-            "Content-Type": "application/json"
-        }
-        payload = {
-            "model": self.model_name,
-            "messages": [{"role": self.role, "content": prompt}]
-        }
-        response = requests.post(self.api_url, headers=headers, data=json.dumps(payload))
-        if response.status_code == 200:
-            return response.json()["choices"][0]["message"]["content"]
-        else:
-            raise Exception(f"Request failed: {response.status_code} - {response.text}")
-
-def call_openrouter(llm, prompt, retries=5, delay=5):
-    for attempt in range(retries):
-        try:
-            return llm.get_response(prompt)
-        except Exception as e:
-            print(f"Retry {attempt+1}: {e}")
-            time.sleep(delay)
-    return "ERROR"
-
-def prepare_claude_ranking_prompt(question, correct, distractors):
+def prepare_ranking_prompt(question, correct, distractors):
     distractor_text = "\n".join([f"Choice {i+1}: {d}" for i, d in enumerate(distractors)])
     return f"""You are a helpful assistant evaluating multiple-choice distractors.
 
@@ -55,7 +21,6 @@ Your task: Rank the distractors from most to least confusing.
 Only respond with one line in this format: 2 > 1 > 3
 """
 
-
 def generate_and_rank(question):
     gemini_prompt = f"""You are an assistant generating multiple-choice questions.
 
@@ -73,6 +38,7 @@ Question: {question}
 """ 
     gemini_response = gemini_model.generate_content(gemini_prompt).text.strip()
 
+    # Parse output
     lines = [line.strip() for line in gemini_response.splitlines() if line.strip()]
     correct = next((line.split(":", 1)[1].strip() for line in lines if line.startswith("Correct Answer")), "N/A")
     distractors = [line.split(":", 1)[1].strip() for line in lines if line.startswith("Distractor")]
@@ -80,20 +46,19 @@ Question: {question}
     if len(distractors) != 3:
         return gemini_response, "ERROR (Invalid distractor count)"
 
-    claude_llm = OpenRouter("anthropic/claude-3.5-sonnet-20240612", os.getenv("OPENROUTER_API_KEY"))
-    claude_prompt = prepare_claude_ranking_prompt(question, correct, distractors)
-    ranking = call_openrouter(claude_llm, claude_prompt)
+    ranking_prompt = prepare_ranking_prompt(question, correct, distractors)
+    ranking_response = gemini_model.generate_content(ranking_prompt).text.strip()
 
-    return gemini_response, ranking.strip()
+    return gemini_response, ranking_response
 
 iface = gr.Interface(
     fn=generate_and_rank,
     inputs=gr.Textbox(label="Enter your MCQ Question"),
     outputs=[
         gr.Textbox(label="Generated Answer and Distractors (Gemini)"),
-        gr.Textbox(label="Claude's Ranking")
+        gr.Textbox(label="Gemini's Ranking of Distractors")
     ],
-    title="Confusing Distractor Generator + Claude Ranking"
+    title="Confusing Distractor Generator + Ranking (Gemini-only)"
 )
 
 iface.launch()