Update app.py

app.py

@@ -1,90 +1,94 @@
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
-import numpy as np
-import cv2
-from PIL import Image
-import pytesseract
-from sentence_transformers import SentenceTransformer, util
-import io
-
-# Define the model name and path to the saved model
-model_path = "E:/grader app/saved_model"  # Replace with the path to your saved model
-
-# Check if CUDA is available, otherwise, fall back to CPU
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using device: {device}")
-
-# Load the tokenizer and model from the local disk
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModelForCausalLM.from_pretrained(
-    model_path, 
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-    device_map="auto" if device == "cuda" else None
-)
-model.to(device)
-
-# Load a smaller version of Sentence-BERT model
-model1 = SentenceTransformer('all-MiniLM-L6-v2')
-
-def get_embedding(text):
-    return model1.encode(text, convert_to_tensor=True)
-
-def calculate_similarity(text1, text2):
-    embedding1 = get_embedding(text1)
-    embedding2 = get_embedding(text2)
-    similarity = util.pytorch_cos_sim(embedding1, embedding2)
-    return similarity.item()
-
-def get_grade(similarity_score):
-    if similarity_score >= 0.9:
-        return 5
-    elif similarity_score >= 0.8:
-        return 4
-    elif similarity_score >= 0.7:
-        return 3
-    elif similarity_score >= 0.6:
-        return 2
-    else:
-        return 1
-
-def extract_text_from_image(image):
-    # Convert PIL image to RGB format
-    image = image.convert('RGB')
-    # Use pytesseract to extract text from the image
-    text = pytesseract.image_to_string(image)
-    return text.strip()
-
-def evaluate_answer(image):
-    student_answer = extract_text_from_image(image)
-    model_answer = "The process of photosynthesis helps plants produce glucose using sunlight."
-    similarity_score = calculate_similarity(student_answer, model_answer)
-    grade = get_grade(similarity_score)
-    feedback = f"Student's answer: {student_answer}\nTeacher's answer: {model_answer}"
-    return grade, similarity_score * 100, feedback
-
-def generate_response(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    
-    # Generate response from the model
-    with torch.no_grad():
-        outputs = model.generate(inputs.input_ids, max_length=150, temperature=0.7)
-    
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
-
-def gradio_interface(image, prompt):
-    grade, similarity_score, feedback = evaluate_answer(image)
-    response = generate_response(prompt)
-    return grade, similarity_score, feedback, response
-
-# Define Gradio interface
-interface = gr.Interface(
-    fn=gradio_interface,
-    inputs=[gr.Image(type="pil"), gr.Textbox(lines=2, placeholder="Enter your prompt here")],
-    outputs=[gr.Label(), gr.Label(), gr.Textbox(), gr.Textbox()],
-    live=True
-)
-
-if __name__ == "__main__":
-    interface.launch()
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+import pytesseract
+from sentence_transformers import SentenceTransformer, util
+import io
+
+model_name = "eachadea/vicuna-7b-1.1"
+
+# Check if CUDA is available, otherwise, fall back to CPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+# Load the tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Load the model
+# If CUDA is available, use float16, otherwise, use float32
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, 
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    device_map="auto" if device == "cuda" else None
+)
+
+# Move model to the appropriate device (CPU or CUDA)
+model.to(device)
+
+# Load a smaller version of Sentence-BERT model
+model1 = SentenceTransformer('all-MiniLM-L6-v2')
+
+def get_embedding(text):
+    return model1.encode(text, convert_to_tensor=True)
+
+def calculate_similarity(text1, text2):
+    embedding1 = get_embedding(text1)
+    embedding2 = get_embedding(text2)
+    similarity = util.pytorch_cos_sim(embedding1, embedding2)
+    return similarity.item()
+
+def get_grade(similarity_score):
+    if similarity_score >= 0.9:
+        return 5
+    elif similarity_score >= 0.8:
+        return 4
+    elif similarity_score >= 0.7:
+        return 3
+    elif similarity_score >= 0.6:
+        return 2
+    else:
+        return 1
+
+def extract_text_from_image(image):
+    # Convert PIL image to RGB format
+    image = image.convert('RGB')
+    # Use pytesseract to extract text from the image
+    text = pytesseract.image_to_string(image)
+    return text.strip()
+
+def evaluate_answer(image):
+    student_answer = extract_text_from_image(image)
+    model_answer = "The process of photosynthesis helps plants produce glucose using sunlight."
+    similarity_score = calculate_similarity(student_answer, model_answer)
+    grade = get_grade(similarity_score)
+    feedback = f"Student's answer: {student_answer}\nTeacher's answer: {model_answer}"
+    return grade, similarity_score * 100, feedback
+
+def generate_response(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    
+    # Generate response from the model
+    with torch.no_grad():
+        outputs = model.generate(inputs.input_ids, max_length=150, temperature=0.7)
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
+
+def gradio_interface(image, prompt):
+    grade, similarity_score, feedback = evaluate_answer(image)
+    response = generate_response(prompt)
+    return grade, similarity_score, feedback, response
+
+# Define Gradio interface
+interface = gr.Interface(
+    fn=gradio_interface,
+    inputs=[gr.Image(type="pil"), gr.Textbox(lines=2, placeholder="Enter your prompt here")],
+    outputs=[gr.Label(), gr.Label(), gr.Textbox(), gr.Textbox()],
+    live=True
+)
+
+if __name__ == "__main__":
+    interface.launch()