Update app.py

app.py

@@ -5,14 +5,15 @@ import torch
 # Load the model and tokenizer
 model_name = 'FridayMaster/fine_tune_embedding'  # Replace with your model's repository name
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)  # Use the appropriate class
 
 # Define a function to generate responses
 def generate_response(prompt):
     inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
     with torch.no_grad():
-        outputs = model.generate(inputs['input_ids'], max_length=150, num_return_sequences=1)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        outputs = model(**inputs)
+    # Customize the response generation as per your model's output
+    response = tokenizer.decode(outputs.logits.argmax(dim=-1), skip_special_tokens=True)
     return response
 
 # Create a Gradio interface
@@ -27,5 +28,3 @@ iface = gr.Interface(
 # Launch the Gradio app
 if __name__ == "__main__":
     iface.launch()
-
-