Update app.py

app.py

@@ -14,6 +14,7 @@ from transformers import (
 )
 import torch
 from typing import Optional
+from huggingface_hub import InferenceClient
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -32,42 +33,26 @@ from tools import (
     query_video
 )
 
+    
 
 class BasicAgent:
-    def __init__(self, model_choice="bertin-gpt-j-6B"):
-        self.model_choice = model_choice
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        
-        if model_choice == "bertin-gpt-j-6B":
-            # Bilingual (Spanish/English) 6B model
-            self.model = AutoModelForCausalLM.from_pretrained(
-                "bertin-project/bertin-gpt-j-6B",
-                torch_dtype=torch.float16,  # Reduce memory usage
-                low_cpu_mem_usage=True
-            ).to(self.device)
-            self.tokenizer = AutoTokenizer.from_pretrained("bertin-project/bertin-gpt-j-6B")
-        elif model_choice == "redpajama-3B":
-            # General-purpose 3B model
-            self.model = AutoModelForCausalLM.from_pretrained(
-                "togethercomputer/RedPajama-INCITE-Chat-3B-v1",
-                torch_dtype=torch.float16
-            ).to(self.device)
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                "togethercomputer/RedPajama-INCITE-Chat-3B-v1"
-            )
-        else:
-            raise ValueError(f"Unsupported model: {model_choice}")
+    def __init__(self):
+        self.model = InferenceClient()
 
-    def __call__(self, question):
-        # GPT-style models (bertin-gpt-j-6B, redpajama-3B)
-        inputs = self.tokenizer(question, return_tensors="pt").to(self.device)
-        outputs = self.model.generate(
-            **inputs,
-            max_new_tokens=100,
-            do_sample=True,
-            temperature=0.7
-        )
-        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+    def __call__(self, question, question_attach):
+        print("Q: "+question="\nattachment: "+question_attach)
+        if question_attach!="":
+            return "attachment analysis is not implemented yet."
+        else:
+            if question.startswith("."):
+                print("it is reversed")
+                question = question[::-1]
+                print("normalized q:"+question)
+            prompt = f"""Answer this GAIA level 1 question:
+                Question: {question}
+                Answer:"""
+            response = self.model.text_generation(prompt, model="mistralai/Mistral-7B-Instruct-v0.1")
+        return response
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
@@ -90,7 +75,7 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
 
     # 1. Instantiate Agent ( modify this part to create your agent)
     try:
-        agent = BasicAgent(model_choice="bertin-gpt-j-6B")
+        agent = BasicAgent()
     except Exception as e:
         print(f"Error instantiating agent: {e}")
         return f"Error initializing agent: {e}", None
@@ -133,7 +118,7 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
         
         submitted_answer = "def_answer"
         try:
-            submitted_answer_raw = agent(question_text)
+            submitted_answer_raw = agent(question_text, question_attach)
             
             if isinstance(submitted_answer_raw, list) and 'generated_text' in submitted_answer_raw[0]:
                 submitted_answer = submitted_answer_raw[0]['generated_text'].strip()