Update app.py

app.py

@@ -1,40 +1,39 @@
 import gradio as gr
 
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
+import torch
 
-MODEL_URL = "https://huggingface.co/serdarakyol/interpress-turkish-news-classification"
-WEBSITE_URL = "https://www.kodiks.com/ai_solutions.html"
+MODEL_URL = "kingabzpro/Llama-3.1-8B-Instruct-Mental-Health-Classification"
 
-tokenizer = AutoTokenizer.from_pretrained("serdarakyol/interpress-turkish-news-classification")
-model = AutoModelForSequenceClassification.from_pretrained("serdarakyol/interpress-turkish-news-classification")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_URL)
+
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_URL, low_cpu_mem_usage=True, return_dict=True,torch_dtype=torch.float16,
+    device_map="cpu")
 
 def prediction(news):
     # create pipeline
-    clasifer = pipeline("sentiment-analysis", tokenizer=tokenizer, model=model, return_all_scores=True)
-
-    preds = clasifer(news)
+    clasifer = pipeline("text-generation", tokenizer=tokenizer, model=model, torch_dtype=torch.float16,
+    device_map="cpu",)
 
-    preds_dict={}
-    for pred in preds[0]:
-        preds_dict[pred['label']] = pred['score']
+    outputs = pipe(prompt, max_new_tokens=2, do_sample=True, temperature=0.1)
+    preds = outputs[0]["generated_text"].split("label: ")[-1].strip())
 
-    return preds_dict
+    return preds
 
 
 gradio_ui = gr.Interface(
     fn=prediction,
-    title="Turkish News Classification",
-    description=f"Enter Turkish news article to see the category of the news.\n For this classification, the {MODEL_URL} model was used.",
+    title="Mental Health Disorder Classification",
+    description=f"Input the text to generate a Mental Health Disorder.\n For this classification, the {MODEL_URL} model was used.",
     examples=[
-        ['Son dakika haberi... Sağlık Bakanlığı corona virüsü (koronavirüs) için son gelişmeleri ve Türkiye Günlük Korona Tablosu\'nu açıkladı. Türkiye\'de son 24 saatte 355 bin 947 Kovid-19 testi yapıldı, 26 bin 896 kişinin testi pozitif çıktı, 210 kişi yaşamını yitirdi. Tabloda vaka ve vefat sayılarına dikkat çeken Bakan Koca, "Bu ölümlerin bir kısmı birden çok sebebin bir araya gelmesine bağlı ve önlenebilir. Örneğin, Covid-19’a kronik hastalıkların eşlik etmesi riski artıran bir faktör. Tam doz aşı bu riski azaltıyor." çağrısında bulundu. Öte yandan ikinci doz aşılama oranı %75\'e çıkan Tunceli, mavi kategoriye yükseldi.'],
-        ["Galatasaray'dan sezon başında ayrılan ve Rayo Vallecano'nun yolunu tutan Radamel Falcao, İspanya'da şov yapmaya devam ediyor. Kolombiyalı yıldızın, Barcelona karşısında attığı gol maça damga vurdu."],
-        ["Instagram, 3 yıl önce kullanıcıların beğenisine sunduğu IGTV özelliğini kaldırıyor. IGTV’nin yerine gelecek yeni özelliğin ise platformun Reels'tan elde ettiği başarıyı pekiştirip içerik üreticileri için kullanışlı bir zemin oluşturması bekleniyor."],
-        ["Kamu bankaları Halkbank, Vakıfbank ve Ziraat Bankası ortak bir açıklama yayınlayarak kredi faiz oranlarında indirim yapıldığını açıkladı. Söz konusu indirimler ticari ve konut kredileri kapsıyor. Tüketici kredisi kapsam dışında bırakıldı."]
+        ['trouble sleeping, confused mind, restless heart. All out of tune'],
+        ["In the quiet hours, even the shadows seem too heavy to bear."],
+        ["Riding a tempest of emotions, where ecstatic highs crash into desolate lows without warning."]
     ],
-    inputs=gr.inputs.Textbox(lines=10, label="Paste some news here"),
-    outputs=gr.outputs.Label(num_top_classes=5, type="auto", label="News categories probabilities"),
+    inputs=gr.inputs.Textbox(lines=10, label="Write the text here"),
+    outputs=gr.outputs.Label(num_top_classes=5, type="auto", label="Mental Health Disorder Category"),
     theme="huggingface",
-    article="<p style='text-align: center'>for our other AI works: <a href='https://www.kodiks.com/ai_solutions.html' target='_blank'>https://www.kodiks.com/ai_solutions.html</a> | <a href='https://twitter.com/KodiksBilisim' target='_blank'>Contact us</a></p>",
+    article="<p style='text-align: center'>Please read the tutorial to fine-tune the Llama 3.1 model on Mental Health Classification <a href='https://www.datacamp.com/tutorial/fine-tuning-llama-3-1' target='_blank'>https://www.datacamp.com/tutorial/fine-tuning-llama-3-1</a></p>",
 )
 
 gradio_ui.launch()