Update app.py

app.py

@@ -1,26 +1,27 @@
 import gradio as gr
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Load the model and tokenizer
+# Load the tokenizer
 model_name = "TuringsSolutions/TechLegalV1"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModel.from_pretrained(model_name)
+
+# Load the model
+# Assuming it's a CausalLM model, you might need to adjust based on your model's architecture
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
 
 # Function to make predictions
 def predict(text):
     inputs = tokenizer(text, return_tensors="pt")
     with torch.no_grad():
-        outputs = model(**inputs)
-    # Assuming we need to extract some specific information from outputs
-    # Modify this part based on your model's output format
-    return outputs.last_hidden_state.mean(dim=1).squeeze().tolist()
+        outputs = model.generate(**inputs)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # Create a Gradio interface
 iface = gr.Interface(
     fn=predict,
     inputs=gr.inputs.Textbox(lines=2, placeholder="Enter text here..."),
-    outputs="json",
+    outputs="text",
     title="Tech Legal Model",
     description="A model for analyzing tech legal documents."
 )