Update agent_t5.py

agent_t5.py

@@ -11,8 +11,10 @@ from transformers import AutoTokenizer, T5ForConditionalGeneration
 from retrieval.retrieval import Retrieval, BM25
 import os, time, torch
 from torch.nn import Softmax
+import requests
 
-
+API_URL = "https://api-inference.huggingface.co/models/CreatorPhan/ViQA-small"
+headers = {"Authorization": "Bearer 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"}
 
 class Agent:
     def __init__(self, args=None) -> None:
@@ -46,6 +48,13 @@ class Agent:
 
         return f"Using file from {doc_path.name}"
     
+    def API_call(self, prompt):
+        response = requests.post(API_URL, headers=headers, json={"inputs": prompt}).json()
+        if isinstance(response, list):
+            return response[0]['generated_text']
+        else:
+            time.sleep(3)
+            return self.API_call(prompt)
 
     def asking(self, question):
         s_query = time.time()