Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-import torch.nn.functional as F
 from peft import (
     LoraConfig,
     PeftModel,
@@ -11,21 +10,29 @@ from peft import (
 model_name = "google/gemma-2-2b-it"
 lora_model_name="Anlam-Lab/gemma-2-2b-it-anlamlab-SA-Chatgpt4mini"
 
-
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu", torch_dtype=torch.float16)
-model = PeftModel.from_pretrained(model, lora_model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
 
-def generate_response(text):
-    example = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>Bir duygu analisti olarak sana verilen metinleri analiz et ve aşağıdaki kategorilerden yalnızca birini seçerek metnin duygu durumunu belirle:Positive,Negative,Neutral<|eot_id|><|start_header_id|>user<|end_header_id|>{text}<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""
-    inputs = tokenizer(example, return_tensors="pt")
+model = PeftModel.from_pretrained(model, lora_model_name)
+def generate_response(input_text):
+    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+    
+    generation_config = {
+        "max_length": 512,
+        "temperature": 0.01,
+        "do_sample": True,
+        "pad_token_id": tokenizer.pad_token_id,
+        "eos_token_id": tokenizer.eos_token_id,
+    }
+    
     with torch.no_grad():
-        model_output = model(**inputs)
-    logits = model_output.logits
-    probabilities = F.softmax(logits, dim=-1)
-    top_probs, top_tokens = torch.topk(probabilities[0, -1, :], k=10)
-    predicted_label = tokenizer.decode(top_tokens[0])
-    return predicted_label
+        outputs = model.generate(
+            **inputs,
+            **generation_config
+        )
+    
+    response = tokenizer.decode(outputs[0])
+    return response.split("<start_of_turn>model\n")[1].split("<end_of_turn>")[0]
 
 iface = gr.Interface(
     fn=generate_response,