Update app.py

app.py

@@ -1,22 +1,29 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+# from transformers import AutoTokenizer, AutoModelForCausalLM
+# import torch
 
-# Load model and tokenizer
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = AutoModelForCausalLM.from_pretrained(
-    "universitytehran/PersianMind-v1.0",
-    torch_dtype=torch.bfloat16,
-)
-tokenizer = AutoTokenizer.from_pretrained("universitytehran/PersianMind-v1.0")
+# # Load model and tokenizer
+# device = "cuda" if torch.cuda.is_available() else "cpu"
+# model = AutoModelForCausalLM.from_pretrained(
+#     "universitytehran/PersianMind-v1.0",
+#     torch_dtype=torch.bfloat16,
+# )
+# tokenizer = AutoTokenizer.from_pretrained("universitytehran/PersianMind-v1.0")
 
-# Conversation template
-TEMPLATE = "{context}\nYou: {prompt}\nPersianMind: "
-CONTEXT = "به انگلیسی ترجمه کن"
-prompt = "به انگلیسی ترجمه کن: وزیر امور خارجه ایران در نشست خبری اعلام کرد که مذاکرات هسته‌ای به دلیل اختلافات اساسی درباره اجرای تعهدات متوقف شده است. او تأکید کرد که ایران همچنان به دنبال یک توافق عادلانه و متوازن است، اما طرف‌های مقابل باید از سیاست‌های دوگانه خود دست بردارند"
+# # Conversation template
+# TEMPLATE = "{context}\nYou: {prompt}\nPersianMind: "
+# CONTEXT = "به انگلیسی ترجمه کن"
+# prompt = "به انگلیسی ترجمه کن: وزیر امور خارجه ایران در نشست خبری اعلام کرد که مذاکرات هسته‌ای به دلیل اختلافات اساسی درباره اجرای تعهدات متوقف شده است. او تأکید کرد که ایران همچنان به دنبال یک توافق عادلانه و متوازن است، اما طرف‌های مقابل باید از سیاست‌های دوگانه خود دست بردارند"
 
-model_input = prompt 
-input_tokens = tokenizer(model_input, return_tensors="pt").to(device)
-generate_ids = model.generate(**input_tokens, max_new_tokens=512, do_sample=True, repetition_penalty=1.1)
-model_output = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-response = model_output[len(model_input):]
-print(response)
+# model_input = prompt 
+# input_tokens = tokenizer(model_input, return_tensors="pt").to(device)
+# generate_ids = model.generate(**input_tokens, max_new_tokens=512, do_sample=True, repetition_penalty=1.1)
+# model_output = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+# response = model_output[len(model_input):]
+# print(response)
+
+
+# Load model directly
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
+model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")