Update main.py

main.py

@@ -1,24 +1,66 @@
-from flask import Flask, request
+from flask import Flask, request, jsonify
 from twilio.twiml.messaging_response import MessagingResponse
-
-sms_chain = LLMChain(
-    llm=Baseten(model="YOUR-MODEL-VERSION-ID"), 
-    prompt=prompt,
-    memory=ConversationBufferWindowMemory(k=2),
-    llm_kwargs={"max_length": 4096}
-)
+from huggingface_hub import InferenceClient
 
 app = Flask(__name__)
 
+# Initialize InferenceClient
+client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.1")
+conversation_history = {}
+
+def format_prompt(message, history):
+    prompt = "<s>"
+    for user_prompt, bot_response in history:
+        prompt += f"[INST] {user_prompt} [/INST]"
+        prompt += f" {bot_response}</s> "
+    prompt += f"[INST] {message} [/INST]"
+    return prompt
+
+def generate(prompt, history, temperature=0.9, max_new_tokens=100, top_p=0.95, repetition_penalty=1.0):
+    formatted_prompt = format_prompt(prompt, history)
+
+    generate_kwargs = dict(
+        temperature=temperature,
+        max_new_tokens=max_new_tokens,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        do_sample=True,
+        seed=42,
+    )
+
+    response = client.text_generation(
+        formatted_prompt,
+        **generate_kwargs,
+        stream=True,
+        details=True,
+        return_full_text=False
+    )
+
+    output = ""
+    for token in response:
+        if hasattr(token, 'token') and hasattr(token.token, 'text'):
+            output += token.token.text
+
+    return output
 
 @app.route("/sms", methods=['GET', 'POST'])
 def sms():
-    resp = MessagingResponse()
-    inb_msg = request.form['Body'].lower().strip()
-    output = sms_chain.predict(sms_input=inb_msg)
-    print(output)
-    resp.message(output)
+    try:
+        resp = MessagingResponse()
+        sender_number = request.form['From']
+        inb_msg = request.form['Body'].strip()
+
+        history = conversation_history.get(sender_number, [])
+        output = generate(inb_msg, history)
+        history.append((inb_msg, output))
+        conversation_history[sender_number] = history
+
+        resp.message(output)
+    except Exception as e:
+        print(f"Error processing SMS: {str(e)}")
+        resp.message("Sorry, an error occurred. Please try again later.")
+    
     return str(resp)
 
 if __name__ == "__main__":
-    app.run(debug=True)
+    app.run(debug=True, port=5000)