Update app.py

app.py

@@ -2,20 +2,30 @@ import os
 import gradio as gr
 import pandas as pd
 import comtradeapicall
-from huggingface_hub import InferenceClient
+import spaces
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from deep_translator import GoogleTranslator
-
+import torch
 
 # کلید COMTRADE
 subscription_key = os.getenv("COMTRADE_API_KEY", "")
 # توکن Hugging Face
 hf_token = os.getenv("HF_API_TOKEN")
 
+# تنظیم کوانت‌سازی
+quantization_config = BitsAndBytesConfig(load_in_4bit=True)
 
-client = InferenceClient(token=hf_token)
-translator = GoogleTranslator(source='en', target='fa')
-
+# بارگذاری توکنایزر و مدل
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b", token=hf_token)
+model = AutoModelForCausalLM.from_pretrained(
+    "google/gemma-7b",
+    token=hf_token,
+    quantization_config=quantization_config,
+    device_map="auto",
+    torch_dtype=torch.float16
+)
 
+# تابع دریافت اطلاعات واردکنندگان
 def get_importers(hs_code: str, year: str, month: str):
     period = f"{year}{int(month):02d}"
     df = comtradeapicall.previewFinalData(
@@ -36,10 +46,12 @@ def get_importers(hs_code: str, year: str, month: str):
     result.columns = ["کد کشور", "نام کشور", "ارزش CIF"]
     return result
 
-
+# تابع ارائه مشاوره با استفاده از GPU
+@spaces.GPU(duration=120)
 def provide_advice(table_data: pd.DataFrame, hs_code: str, year: str, month: str):
     if table_data is None or table_data.empty:
         return "ابتدا باید اطلاعات واردات را نمایش دهید."
+    
     table_str = table_data.to_string(index=False)
     period = f"{year}/{int(month):02d}"
     prompt = (
@@ -49,19 +61,21 @@ def provide_advice(table_data: pd.DataFrame, hs_code: str, year: str, month: str
     )
     print("پرامپت ساخته‌شده:")
     print(prompt)
+
     try:
-        print("در حال فراخوانی مدل mistralai/Mixtral-8x7B-Instruct-v0.1...")
-        outputs = client.text_generation(
-            prompt=prompt,
-            model="mistralai/Mixtral-8x7B-Instruct-v0.1",
-            max_new_tokens=1024  # افزایش برای تکمیل جملات
+        input_ids = tokenizer(prompt, return_tensors="pt").to("cuda")
+        outputs = model.generate(
+            **input_ids,
+            max_new_tokens=1024,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9
         )
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         print("خروجی مدل دریافت شد (به انگلیسی):")
-        print(outputs)
-
+        print(generated_text)
 
-        # ترجمه خروجی به فارسی
-        translated_outputs = translator.translate(outputs)
+        translated_outputs = translator.translate(generated_text)
         print("خروجی ترجمه‌شده به فارسی:")
         print(translated_outputs)
         return translated_outputs
@@ -70,12 +84,11 @@ def provide_advice(table_data: pd.DataFrame, hs_code: str, year: str, month: str
         print(error_msg)
         return error_msg
 
-
+# تنظیمات رابط Gradio
 current_year = pd.Timestamp.now().year
 years = [str(y) for y in range(2000, current_year+1)]
 months = [str(m) for m in range(1, 13)]
 
-
 with gr.Blocks() as demo:
     gr.Markdown("##تولید شده توسط DIGINORON نمایش کشورهایی که یک کالا را وارد کرده‌اند")
     with gr.Row():
@@ -90,17 +103,14 @@ with gr.Blocks() as demo:
     )
     btn_show.click(get_importers, [inp_hs, inp_year, inp_month], out_table)
 
-
     btn_advice = gr.Button("ارائه مشاوره تخصصی")
     out_advice = gr.Textbox(label="مشاوره تخصصی", lines=6)
 
-
     btn_advice.click(
         provide_advice,
         inputs=[out_table, inp_hs, inp_year, inp_month],
         outputs=out_advice
     )
 
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()