up

app.py

@@ -6,12 +6,20 @@ import docx2txt
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
 import difflib
-from predict import run_prediction
+from huggingface_hub import InferenceApi  # Import Hugging Face API
 
-# ========== CONFIG ==========
+# ========== CONFIG ========== 
 st.set_page_config(page_title="📑 Contract Analyzer", layout="wide")
 
 # ========== FUNCTIONS ==========
+
+# Tải mô hình Hugging Face từ Hub
+@st.cache_resource
+def load_inference_api():
+    return InferenceApi(repo_id="HuggingFaceH4/zephyr-7b-beta")  # Mô hình Zephyr
+
+inference_api = load_inference_api()
+
 def extract_text_from_pdf(uploaded_file):
     try:
         with pdfplumber.open(uploaded_file) as pdf:
@@ -66,7 +74,17 @@ def compute_similarity(text1, text2):
     except:
         return difflib.SequenceMatcher(None, text1, text2).ratio() * 100
 
-# ========== MAIN ==========
+# Hàm truy vấn Zephyr từ Hugging Face
+def query_zephyr_model(text1, text2, question):
+    prompt = f"Compare the following two contracts and answer the question:\nText 1: {text1}\nText 2: {text2}\nQuestion: {question}"
+    try:
+        result = inference_api(inputs=prompt)
+        return result['generated_text']
+    except Exception as e:
+        st.error(f"Error querying the model: {e}")
+        return None
+
+# ========== MAIN ========== 
 def main():
     st.title("📑 Contract Analyzer")
     st.markdown("Upload two contracts, compare them, and ask any question!")
@@ -116,8 +134,8 @@ def main():
             st.subheader("Answer from Document 1")
             with st.spinner("Analyzing..."):
                 try:
-                    pred1 = run_prediction([user_question], text1, model_name='marshmellow77/roberta-base-cuad', n_best_size=5)
-                    st.success(pred1.get('0', 'No answer found'))
+                    pred1 = query_zephyr_model(text1, text2, user_question)
+                    st.success(pred1)
                 except Exception as e:
                     st.error(f"Failed on Document 1: {e}")
 
@@ -125,8 +143,8 @@ def main():
             st.subheader("Answer from Document 2")
             with st.spinner("Analyzing..."):
                 try:
-                    pred2 = run_prediction([user_question], text2, model_name='marshmellow77/roberta-base-cuad', n_best_size=5)
-                    st.success(pred2.get('0', 'No answer found'))
+                    pred2 = query_zephyr_model(text1, text2, user_question)
+                    st.success(pred2)
                 except Exception as e:
                     st.error(f"Failed on Document 2: {e}")